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()
|