#!/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()