pic_npz copy.py

#!/usr/bin/env python3
"""
将文件夹下所有PNG或JPG文件读取并生成对应NPZ文件
基于 sample_ddp_new.py 中的 create_npz_from_sample_folder 函数改进
支持自动检测图片数量，支持PNG和JPG格式，输出到父级目录
支持从 metadata.jsonl 文件读取图片路径
"""

import os
import argparse
import numpy as np
from PIL import Image
from tqdm import tqdm
import glob
import json


def create_npz_from_metadata(metadata_jsonl_path, output_path=None):
    """
    从 metadata.jsonl 文件读取图片路径并构建 .npz 文件
    
    Args:
        metadata_jsonl_path (str): metadata.jsonl 文件路径
        output_path (str, optional): 输出 npz 文件路径，默认在 metadata.jsonl 同目录下生成
        
    Returns:
        str: 生成的 npz 文件路径
    """
    # 确保 metadata.jsonl 存在
    if not os.path.exists(metadata_jsonl_path):
        raise ValueError(f"metadata.jsonl 文件不存在: {metadata_jsonl_path}")
    
    # 获取基础目录
    base_dir = os.path.dirname(metadata_jsonl_path)
    
    # 读取 metadata.jsonl
    image_files = []
    with open(metadata_jsonl_path, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if line:
                try:
                    data = json.loads(line)
                    file_name = data.get('file_name')
                    if file_name:
                        full_path = os.path.join(base_dir, file_name)
                        image_files.append(full_path)
                except json.JSONDecodeError as e:
                    print(f"警告: 跳过无效的 JSON 行: {e}")
                    continue
    
    if len(image_files) == 0:
        raise ValueError(f"在 {metadata_jsonl_path} 中未找到任何有效的图片路径")
    
    print(f"从 metadata.jsonl 读取到 {len(image_files)} 张图片路径")
    
    # 读取所有图片
    samples = []
    for img_path in tqdm(image_files, desc="读取图片并转换为numpy数组"):
        try:
            # 打开图片并转换为RGB格式（确保一致性）
            with Image.open(img_path) as img:
                # 转换为RGB，确保所有图片都是3通道
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # 将图片resize到512x512
                img = img.resize((512, 512), Image.LANCZOS)
                
                sample_np = np.asarray(img).astype(np.uint8)
                
                # 确保图片是3通道
                if len(sample_np.shape) != 3 or sample_np.shape[2] != 3:
                    print(f"警告: 跳过非3通道图片 {img_path}, 形状: {sample_np.shape}")
                    continue
                    
                samples.append(sample_np)
                
        except Exception as e:
            print(f"警告: 无法读取图片 {img_path}: {e}")
            continue
    
    if len(samples) == 0:
        raise ValueError("没有成功读取任何有效的图片文件")
    
    # 转换为numpy数组
    samples = np.stack(samples)
    print(f"成功读取 {len(samples)} 张图片，形状: {samples.shape}")
    
    # 验证数据形状
    assert len(samples.shape) == 4, f"期望4维数组，得到形状: {samples.shape}"
    assert samples.shape[3] == 3, f"期望3通道图片，得到: {samples.shape[3]}通道"
    
    # 生成输出路径
    if output_path is None:
        base_name = os.path.splitext(os.path.basename(metadata_jsonl_path))[0]
        output_path = os.path.join(base_dir, f"{base_name}.npz")
    
    # 保存为npz文件
    np.savez(output_path, arr_0=samples)
    print(f"已保存 .npz 文件到 {output_path} [形状={samples.shape}]")
    
    return output_path


def main():
    """
    主函数：解析命令行参数并执行图片到npz的转换
    """
    parser = argparse.ArgumentParser(
        description="将文件夹下所有PNG或JPG文件转换为NPZ格式",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  python pic_npz.py /path/to/image/folder
  python pic_npz.py /path/to/image/folder --output-dir /custom/output/path
        """
    )
    
    parser.add_argument(
        "--image_folder", 
        type=str, 
        default="/gemini/space/gzy_new/models/Sida/sd3_rectified_samples",
        help="包含PNG或JPG图片文件的文件夹路径"
    )
    
    # parser.add_argument(
    #     "--metadata_jsonl",
    #     type=str,
    #     default="/gemini/space/hsd/project/dataset/cc3m-wds/validation/metadata.jsonl",
    #     help="metadata.jsonl 文件路径，用于从 JSONL 文件读取图片路径"
    # )
    
    parser.add_argument(
        "--output-dir", 
        type=str, 
        default=None,
        help="自定义输出目录（默认为输入文件夹的父级目录或 metadata.jsonl 所在目录）"
    )
    
    args = parser.parse_args()
    
    try:
        if args.metadata_jsonl and os.path.exists(args.metadata_jsonl):
            # 使用 metadata.jsonl
            metadata_path = os.path.abspath(args.metadata_jsonl)
            base_dir = os.path.dirname(metadata_path)
            base_name = os.path.splitext(os.path.basename(metadata_path))[0]
            
            if args.output_dir:
                os.makedirs(args.output_dir, exist_ok=True)
                output_path = os.path.join(args.output_dir, f"{base_name}.npz")
            else:
                output_path = os.path.join(base_dir, f"{base_name}.npz")
            
            npz_path = create_npz_from_metadata(metadata_path, output_path)
        else:
            # 使用图片文件夹
            image_folder_path = os.path.abspath(args.image_folder)
            
            if args.output_dir:
                # 如果指定了输出目录，修改生成逻辑
                folder_name = os.path.basename(image_folder_path.rstrip('/'))
                custom_output_path = os.path.join(args.output_dir, f"{folder_name}.npz")
                
                # 创建输出目录（如果不存在）
                os.makedirs(args.output_dir, exist_ok=True)
                
                # 临时修改函数以支持自定义输出路径
                npz_path = create_npz_from_image_folder_custom(image_folder_path, custom_output_path)
            else:
                npz_path = create_npz_from_image_folder(image_folder_path)
        
        print(f"转换完成！NPZ文件已保存至: {npz_path}")
        
    except Exception as e:
        print(f"错误: {e}")
        return 1
    
    return 0


def create_npz_from_image_folder_custom(image_folder_path, output_path):
    """
    从包含图片的文件夹构建单个 .npz 文件（自定义输出路径版本）
    
    Args:
        image_folder_path (str): 包含图片文件的文件夹路径
        output_path (str): 输出npz文件的完整路径
        
    Returns:
        str: 生成的 npz 文件路径
    """
    # 确保路径存在
    if not os.path.exists(image_folder_path):
        raise ValueError(f"文件夹路径不存在: {image_folder_path}")
    
    # 获取所有支持的图片文件
    supported_extensions = ['*.png', '*.PNG', '*.jpg', '*.JPG', '*.jpeg', '*.JPEG']
    image_files = []
    
    for extension in supported_extensions:
        pattern = os.path.join(image_folder_path, extension)
        image_files.extend(glob.glob(pattern))
    
    # 按文件名排序确保一致性
    image_files.sort()
    
    if len(image_files) == 0:
        raise ValueError(f"在文件夹 {image_folder_path} 中未找到任何PNG或JPG图片文件")
    
    print(f"找到 {len(image_files)} 张图片文件")
    
    # 读取所有图片
    samples = []
    for img_path in tqdm(image_files, desc="读取图片并转换为numpy数组"):
        try:
            # 打开图片并转换为RGB格式（确保一致性）
            with Image.open(img_path) as img:
                # 转换为RGB，确保所有图片都是3通道
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # 将图片resize到512x512
                img = img.resize((512, 512), Image.LANCZOS)
                
                sample_np = np.asarray(img).astype(np.uint8)
                
                # 确保图片是3通道
                if len(sample_np.shape) != 3 or sample_np.shape[2] != 3:
                    print(f"警告: 跳过非3通道图片 {img_path}, 形状: {sample_np.shape}")
                    continue
                    
                samples.append(sample_np)
                
        except Exception as e:
            print(f"警告: 无法读取图片 {img_path}: {e}")
            continue
    
    if len(samples) == 0:
        raise ValueError("没有成功读取任何有效的图片文件")
    
    # 转换为numpy数组
    samples = np.stack(samples)
    print(f"成功读取 {len(samples)} 张图片，形状: {samples.shape}")
    
    # 验证数据形状
    assert len(samples.shape) == 4, f"期望4维数组，得到形状: {samples.shape}"
    assert samples.shape[3] == 3, f"期望3通道图片，得到: {samples.shape[3]}通道"
    
    # 保存为npz文件
    np.savez(output_path, arr_0=samples)
    print(f"已保存 .npz 文件到 {output_path} [形状={samples.shape}]")
    
    return output_path


if __name__ == "__main__":
    exit(main())