【异常错误】不论在测试还是训练的时候都要归一化并反归一化，保存图像方法要合适

一、图像归一化问题

这里没有在测试的时候给图像归一化，会导致输出的图像颜色上有很大问题，例如就会变成这样，

所以需要在进入model之前归一化，并且在输出之后反归一化才可以。

二、在保存图像的时候

使用“vutils.save_image”而别使用ToPILImage()，否则也会导致你的图像出现很大的问题

有问题：

    import torchvision.utils as vutils
    
    # 方式一：正常保存，没有问题
    vutils.save_image(decoded_img, os.path.join(out_dir_path, filename))
 

没有问题： 

    # 方式二：使用PIL保存，有问题
    out_img = ToPILImage()(decoded_img)
    out_img.save(os.path.join(out_dir_path, filename))

正确代码示例：

来自于：https://github.com/CompVis/latent-diffusion/pull/353/files#diff-d7c29e5cdd97960f10932618a1693a094a186fddb09b5d781d39d49508a6b6a7

import torch
from ldm.util import instantiate_from_config
from omegaconf import OmegaConf
from PIL import Image
import torchvision.transforms as T
import os
import torchvision.utils as vutils
import argparse
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
def load_model_from_config(config, ckpt):
    print(f"Loading model from {ckpt}")
    pl_sd = torch.load(ckpt, map_location="cpu")
    global_step = pl_sd["global_step"]
    sd = pl_sd["state_dict"]
    model = instantiate_from_config(config.model)
    m, u = model.load_state_dict(sd, strict=False)
    model.to(device)
    model.eval()
    return {"model": model}, global_step
 
 
def load_and_preprocess_image(image_path, resize_shape=(256, 256)):
    transform = T.Compose([
        T.Resize(resize_shape),
        T.ToTensor(),
        T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
    image = Image.open(image_path).convert("RGB")
    return transform(image).unsqueeze(0).to(device)
 
 
def reconstruct_image(model, image_tensor):
    with torch.no_grad():
        reconstructed_img, _ = model(image_tensor)
        return reconstructed_img
 
 
def save_image(tensor, filename):
    print("Tensor Type:", type(tensor))  # Debugging line to confirm tensor type
    if isinstance(tensor, torch.Tensor):
        tensor = (tensor + 1) / 2  # Normalize if the tensor is in the range [-1, 1]
        vutils.save_image(tensor, filename)
    else:
        print("The input is not a tensor.")
 
 
def reconstruct_and_save_images(input_dir, output_dir, model):
    for image_name in os.listdir(input_dir):
        if not image_name.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
            continue
        # 1、加载图像
        image_path = os.path.join(input_dir, image_name)
        # 2、将图像转为Tensor
        image_tensor = load_and_preprocess_image(image_path)
        # 3、将图像Tensor经过model重建为tensor
        reconstructed_img = reconstruct_image(model, image_tensor)
        # 4、保存重建的图像
        output_path = os.path.join(output_dir, image_name)
        save_image(reconstructed_img, output_path)
 
 
def main(config_path, ckpt_path, input_dir, output_dir):
    config = OmegaConf.load(config_path)
    model_info, step = load_model_from_config(config, ckpt_path)
    model = model_info["model"]
 
    os.makedirs(output_dir, exist_ok=True)
    reconstruct_and_save_images(input_dir, output_dir, model)
 
 
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Reconstruct images from training autoencoder models")
    parser.add_argument('--config', type=str, required=True, help='Path to model config YAML file')
    parser.add_argument('--ckpt', type=str, required=True, help='Path to model checkpoint file')
    parser.add_argument('--input_dir', type=str, required=True, help='Directory where input images are stored')
    parser.add_argument('--output_dir', type=str, required=True, help='Directory where output images will be saved')
 
    args = parser.parse_args()
    main(args.config, args.ckpt, args.input_dir, args.output_dir)