File size: 2,652 Bytes
a6dd040 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import os
import torch
import torchvision.utils as vutils
from einops import rearrange
from torchvision.transforms.functional import to_pil_image
import numpy as np
from PIL import Image
import cv2
def save_output_images(
output_color: torch.Tensor,
save_dir: str = "./output_images",
prefix: str = "output"
):
"""
保存 output.color 张量中的所有图像为 PNG 文件。
Args:
output_color (Tensor): shape = [b, v, c, h, w]
save_dir (str): 保存图片的目录
prefix (str): 文件名前缀
"""
os.makedirs(save_dir, exist_ok=True)
# 将 [b, v, c, h, w] 变为 [(b * v), c, h, w]
images = rearrange(output_color, "b v c h w -> (b v) c h w")
# 遍历每张图像保存
for i, img_tensor in enumerate(images):
# 确保数值范围在 [0, 1] 之间
img_tensor = img_tensor.clamp(0, 1)
# 转为 PIL 图像，需调整通道顺序为 HWC，并且转换为 uint8 类型
img_tensor = img_tensor.permute(1, 2, 0) # 从 (C, H, W) 转为 (H, W, C)
# 使用 .detach() 来切断计算图，防止在转换时出现问题
img_pil = Image.fromarray((img_tensor.detach().cpu().numpy() * 255).astype("uint8"))
# 保存图像
img_pil.save(os.path.join(save_dir, f"{prefix}_{i:04d}.png"))
def save_depth_images(depth_tensor: torch.Tensor,
out_dir: str,
prefix: str = "depth",
cmap: int = cv2.COLORMAP_INFERNO):
"""
保存深度图为伪彩色 PNG。
Args:
depth_tensor (torch.Tensor): shape [B, V, H, W], float32
out_dir (str): 保存目录，会自动创建
prefix (str): 文件名前缀, 最终文件名会是 prefix_b{b}_v{v}.png
cmap (int): OpenCV 调色板，如 COLORMAP_INFERNO
"""
os.makedirs(out_dir, exist_ok=True)
depth_np = depth_tensor.detach().cpu().numpy()
B, V, H, W = depth_np.shape
for b in range(B):
for v in range(V):
d = depth_np[b, v]
# 归一化到 [0,255]
d_min, d_max = float(d.min()), float(d.max())
if d_max > d_min:
d_norm = (d - d_min) / (d_max - d_min)
else:
d_norm = np.zeros_like(d)
d_8u = (d_norm * 255).astype(np.uint8)
# 伪彩色
colored = cv2.applyColorMap(d_8u, cmap)[:, :, ::-1] # BGR->RGB
# 保存
fname = f"{prefix}_b{b}_v{v}.png"
path = os.path.join(out_dir, fname)
Image.fromarray(colored).save(path)
print(f"Saved {B*V} depth images to {out_dir}")
|