TaxonomyProject/SimulationImage

1"""
2Generate seenable_obj_dict.json for all scenes.
3Example:
4python code/generate_seenable_object_dict.py /home/xwang378/scratch/2025/Taxonomy/Data/simulationImage/ --scene-workers 8 --camera-workers 8
5"""
6
7import os
8import json
9import argparse
10import numpy as np
11from PIL import Image
12from concurrent.futures import ProcessPoolExecutor, as_completed
13from multiprocessing import cpu_count
14
15
16def process_camera(save_path, camera):
17    """处理单个相机的数据"""
18    image_dir = os.path.join(save_path, camera)
19    seg_file = os.path.join(image_dir, "seg.png")
20    obj_anno_file = os.path.join(image_dir, "object_annots.json")
21
22    if not os.path.exists(seg_file) or not os.path.exists(obj_anno_file):
23        return f"[Warning] Missing files in {camera}, skipped.", False
24    
25    if os.path.exists(os.path.join(save_path, camera, "seenable_obj_dict.json")):
26        return f"[Warning] seenable_obj_dict.json already exists for {save_path.split('/')[-1]}/{camera}, skipped.", False
27    
28    # 读取 segmentation 图像和标注
29    with open(obj_anno_file, "r") as f:
30        obj_anno = json.load(f)
31    obj_annos = obj_anno.get("outputs", [])
32
33    seg = np.array(Image.open(seg_file))
34
35    rgb_mask = seg[:, :, :3]
36    
37    # ⚡ 只获取唯一颜色，不计数
38    unique_colors = np.unique(rgb_mask.reshape(-1, 3), axis=0)
39    color_set = set(map(tuple, unique_colors))
40    
41    obj_dict = {
42        obj_anno["object_id"]: tuple(obj_anno["color"][0:3])
43        for obj_anno in obj_annos
44        if tuple(obj_anno["color"][0:3]) in color_set
45    }
46    
47    output_file = os.path.join(save_path, camera, "seenable_obj_dict.json")
48    with open(output_file, "w") as f:
49        json.dump(obj_dict, f, indent=4)
50    
51    return f"[Saved] {output_file}", True
52
53
54def process_scene(image_dir, scene_name, max_workers=None):
55    save_path = os.path.join(image_dir, scene_name)
56
57    if not os.path.exists(save_path):
58        print(f"[Error] Scene path not found: {save_path}")
59        return
60
61    camera_list = [x for x in os.listdir(save_path) if not x.endswith(".json")]
62
63    print(f"Processing scene: {scene_name}")
64    print(f"Found {len(camera_list)} camera folders.")
65
66    if not camera_list:
67        print(f"✅ Done processing scene: {scene_name} (no cameras found)\n")
68        return
69
70    # 并行处理所有相机
71    success_count = 0
72    with ProcessPoolExecutor(max_workers=max_workers) as executor:
73        # 提交所有任务
74        futures = {
75            executor.submit(process_camera, save_path, camera): camera 
76            for camera in camera_list
77        }
78        
79        # 收集结果
80        for future in as_completed(futures):
81            camera = futures[future]
82            try:
83                message, success = future.result()
84                print(message)
85                if success:
86                    success_count += 1
87            except Exception as exc:
88                print(f"[Error] {camera} generated an exception: {exc}")
89
90    print(f"✅ Done processing scene: {scene_name} ({success_count}/{len(camera_list)} cameras processed)\n")
91
92
93if __name__ == "__main__":
94    parser = argparse.ArgumentParser(description="Generate seenable_obj_dict.json for all scenes.")
95    parser.add_argument("image_dir", type=str, help="Directory containing the image folders")
96    parser.add_argument("--scene-workers", type=int, default=None, 
97                        help="Number of parallel workers for scene-level processing (default: CPU count)")
98    parser.add_argument("--camera-workers", type=int, default=None,
99                        help="Number of parallel workers for camera-level processing (default: CPU count)")
100    args = parser.parse_args()
101    
102    # 获取所有场景
103    batch_dir = ['zehan', 'placement', 'jiawei', 'luoxin', 'additional']
104    scenes = []
105    for batch in batch_dir:
106        for scene in os.listdir(os.path.join(args.image_dir, batch)):
107            if os.path.isdir(os.path.join(args.image_dir, batch, scene)):
108                scenes.append(os.path.join(batch, scene))
109    
110    print(f"Found {len(scenes)} scenes to process.")
111    print(f"Scene-level workers: {args.scene_workers or cpu_count()}")
112    print(f"Camera-level workers: {args.camera_workers or cpu_count()}\n")
113    
114    # 并行处理所有场景
115    with ProcessPoolExecutor(max_workers=args.scene_workers) as executor:
116        futures = {
117            executor.submit(process_scene, args.image_dir, scene, args.camera_workers): scene 
118            for scene in scenes
119        }
120        
121        for future in as_completed(futures):
122            scene = futures[future]
123            try:
124                future.result()
125            except Exception as exc:
126                print(f"[Error] Scene {scene} generated an exception: {exc}")
127    
128    print("\n🎉 All scenes processed!")