npy转png

npy转png

import numpy as np
import os
from PIL import Image

def npy_to_png(npy_path, png_path):
    data_npy = np.load(npy_path)
    # 这里根据需要进行转换
    # 归一化到0-255范围， 并转换为uint8类型
    data_npy = np.clip(
      (data_npy - data_npy.min()) / (data_npy.max() - data_npy.min()) * 255, 
      0, 255
    ).astype(np.uint8)

    # 确保数组是三维的，形如(height, width, channels)
    # 如果是灰度图像，则应该是二维数组(height, width)
    if data_npy.ndim == 2:
        # 灰度图像
        image = Image.fromarray(data_npy)
    elif data_npy.ndim == 3:
        # 彩色图像
        if data_npy.shape[2] not in [3, 4]:
            raise ValueError("数组必须是(H, W, 3)或(H, W, 4)格式")
        image = Image.fromarray(data_npy)
    else:
        raise ValueError("数组维度必须为2或3")

    image.save(png_path)

if __name__ == '__main__':
  npy_dir_path = "./data/npy"
  npy_name_list = os.listdir(npy_dir_path)
  for npy_name in npy_name_list:
    npy_path = os.path.join(npy_dir_path, npy_name)
    png_path = os.path.join("./data/png", npy_name.replace(".npy", ".png"))
    npy_to_png(npy_path, png_path)