import os  
import json  

def collect_files(data_dir):  
    """  
    遍历 data 文件夹下的各子文件夹，以文件名的前 7 个字符为键，将对应的文件路径整理为字典  
    """  
    file_dict = {}  
    
    # 子文件夹列表  
    subfolders = ['densepose', 'videos', 'cloth', 'cloth_mask', 'agnostic_mask', 'agnostic']  
    
    for subfolder in subfolders:  
        subfolder_path = os.path.join(data_dir, subfolder)  
        
        if not os.path.exists(subfolder_path):  
            print(f"Warning: {subfolder_path} 路径不存在")  
            continue  
        
        # 遍历子文件夹中的文件  
        for file_name in os.listdir(subfolder_path):  
            # 只取文件名前 7 个字符用于匹配  
            key = file_name[:7]  
            if key not in file_dict:  
                # 初始化字典键为前 7 个字符的键名  
                file_dict[key] = {}  
            
            # 将当前文件路径保存在子文件夹名称对应的 key 下  
            file_dict[key][subfolder] = os.path.join(subfolder_path, file_name)  
    
    return file_dict  

def generate_json(data_dir, output_file):  
    """  
    生成 JSON 文件，将文件匹配结果输出  
    """  
    files = collect_files(data_dir)  
    result = []  

    # 构造符合格式的 JSON 列表  
    for key, paths in files.items():  
        result.append({  
            "densepose": paths.get("densepose", ""),         # 如果某个字段不存在，则填补为空值  
            "videos": paths.get("videos", ""),  
            "cloth": paths.get("cloth", ""),  
            "cloth_mask": paths.get("cloth_mask", ""),  
            "agnostic_mask": paths.get("agnostic_mask", ""),  
            "agnostic": paths.get("agnostic", "")  
        })  
    
    # 写入到指定路径的 JSON 文件  
    with open(output_file, "w", encoding="utf-8") as f:  
        json.dump(result, f, indent=4, ensure_ascii=False)  
    
    print(f"JSON 文件已生成: {output_file}")  

if __name__ == "__main__":  
    # 要匹配的 data 文件夹路径  
    data_dir = "/mnt/lpai-dione/ssai/cvg/team/wjj/ViViD/data"  
    # 输出的 JSON 文件路径  
    output_file = "/mnt/lpai-dione/ssai/cvg/team/wjj/ViViD/data/vividfuxian_stage1.json"  
    
    generate_json(data_dir, output_file)