108 lines
3.5 KiB
Python
108 lines
3.5 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
from PIL import Image
|
|||
|
|
import os
|
|||
|
|
import sys
|
|||
|
|
|
|||
|
|
def npy_to_images(npy_file):
|
|||
|
|
"""
|
|||
|
|
将NumPy文件转换为图片
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
npy_file: .npy文件路径
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
# 检查文件是否存在
|
|||
|
|
if not os.path.exists(npy_file):
|
|||
|
|
print(f"错误: 文件不存在 - {npy_file}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# 加载NumPy数组
|
|||
|
|
print(f"加载文件: {npy_file}")
|
|||
|
|
arrays = np.load(npy_file)
|
|||
|
|
print(f"数组形状: {arrays.shape}")
|
|||
|
|
print(f"数据类型: {arrays.dtype}")
|
|||
|
|
|
|||
|
|
# 创建输出目录
|
|||
|
|
output_dir = f"images_{os.path.splitext(os.path.basename(npy_file))[0]}"
|
|||
|
|
os.makedirs(output_dir, exist_ok=True)
|
|||
|
|
|
|||
|
|
# 数据类型转换
|
|||
|
|
if arrays.dtype == np.float32 or arrays.dtype == np.float64:
|
|||
|
|
if arrays.max() <= 1.0:
|
|||
|
|
# 范围 [0, 1]，转换为 [0, 255]
|
|||
|
|
arrays = (arrays * 255).astype(np.uint8)
|
|||
|
|
else:
|
|||
|
|
# 范围可能是 [0, 255] 的浮点数
|
|||
|
|
arrays = np.clip(arrays, 0, 255).astype(np.uint8)
|
|||
|
|
else:
|
|||
|
|
arrays = arrays.astype(np.uint8)
|
|||
|
|
|
|||
|
|
# 根据数组维度处理
|
|||
|
|
if len(arrays.shape) == 4:
|
|||
|
|
# 4D数组: 多张图片
|
|||
|
|
print(f"检测到4D数组，转换 {arrays.shape[0]} 张图片")
|
|||
|
|
|
|||
|
|
for i in range(arrays.shape[0]):
|
|||
|
|
frame = arrays[i]
|
|||
|
|
|
|||
|
|
# 处理通道顺序 (CHW -> HWC)
|
|||
|
|
if frame.shape[0] == 3 or frame.shape[0] == 1:
|
|||
|
|
frame = np.transpose(frame, (1, 2, 0))
|
|||
|
|
|
|||
|
|
# 保存图片
|
|||
|
|
if len(frame.shape) == 3 and frame.shape[2] == 1:
|
|||
|
|
frame = frame.squeeze(axis=2) # 移除单通道维度
|
|||
|
|
|
|||
|
|
img = Image.fromarray(frame)
|
|||
|
|
filename = f"{output_dir}/image_{i+1:04d}.jpg"
|
|||
|
|
img.save(filename)
|
|||
|
|
|
|||
|
|
if (i + 1) % 10 == 0 or i == 0:
|
|||
|
|
print(f"已保存: {filename}")
|
|||
|
|
|
|||
|
|
elif len(arrays.shape) == 3:
|
|||
|
|
# 3D数组: 单张图片或多张灰度图
|
|||
|
|
if arrays.shape[2] == 3 or arrays.shape[2] == 1:
|
|||
|
|
# 单张彩色或灰度图片 (H, W, C)
|
|||
|
|
print("检测到3D数组，转换为单张图片")
|
|||
|
|
if arrays.shape[2] == 1:
|
|||
|
|
arrays = arrays.squeeze(axis=2)
|
|||
|
|
img = Image.fromarray(arrays)
|
|||
|
|
filename = f"{output_dir}/image.jpg"
|
|||
|
|
img.save(filename)
|
|||
|
|
print(f"已保存: {filename}")
|
|||
|
|
else:
|
|||
|
|
# 多张灰度图片 (N, H, W)
|
|||
|
|
print(f"检测到3D数组，转换 {arrays.shape[0]} 张灰度图片")
|
|||
|
|
for i in range(arrays.shape[0]):
|
|||
|
|
img = Image.fromarray(arrays[i])
|
|||
|
|
filename = f"{output_dir}/image_{i+1:04d}.jpg"
|
|||
|
|
img.save(filename)
|
|||
|
|
if (i + 1) % 10 == 0 or i == 0:
|
|||
|
|
print(f"已保存: {filename}")
|
|||
|
|
|
|||
|
|
elif len(arrays.shape) == 2:
|
|||
|
|
# 2D数组: 单张灰度图片
|
|||
|
|
print("检测到2D数组，转换为灰度图片")
|
|||
|
|
img = Image.fromarray(arrays)
|
|||
|
|
filename = f"{output_dir}/image.jpg"
|
|||
|
|
img.save(filename)
|
|||
|
|
print(f"已保存: {filename}")
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
print(f"不支持的数组维度: {arrays.shape}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
print(f"所有图片已保存到目录: {output_dir}")
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
"""主函数"""
|
|||
|
|
npy_file = "/root/autodl-tmp/hot_video_analyse/source/scene_change/scene_change_arrays.npy"
|
|||
|
|
npy_to_images(npy_file)
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|