1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176
| import omni.replicator.core as rep import omni.usd import carb import numpy as np
class CabinSyntheticDataGenerator: """ 座舱合成数据生成器 生成内容: 1. 驾驶员姿态(疲劳、分心) 2. 光照变化(白天、夜间、隧道) 3. 遮挡场景(墨镜、口罩) 4. 相机视角变化 """ def __init__(self, usd_stage_path: str): self.stage = omni.usd.get_context().open_stage(usd_stage_path) self.driver = self.stage.GetPrimAtPath('/World/Driver') self.cabin = self.stage.GetPrimAtPath('/World/Cabin') self.camera = self.stage.GetPrimAtPath('/World/Camera') self.randomizers = {} def setup_randomization(self): """设置随机化参数""" self.randomizers['lighting'] = rep.randomizer.lighting( light_types=['DomeLight', 'RectLight', 'SphereLight'], intensity_range=(100, 2000), color_temperature_range=(3000, 7000), ) self.randomizers['material'] = rep.randomizer.material( materials=['seat_leather', 'seat_fabric', 'dash_plastic'], texture_randomization=True, ) self.randomizers['pose'] = rep.randomizer.transform( prims=['/World/Driver/Head', '/World/Driver/Body'], translations=[(-0.5, 0.8, 0.2), (0.5, 1.2, 0.4)], rotations=[(-30, -30, -30), (30, 30, 30)], ) def generate_fatigue_scenarios(self, num_samples: int = 1000): """ 生成疲劳场景数据 疲劳特征: - 眼睑下垂 - 眨眼频率降低 - 头部下垂 """ scenario_config = { 'eye_openness': (0.2, 0.5), 'blink_rate': (5, 10), 'head_pitch': (10, 30), } for i in range(num_samples): self._set_fatigue_params(scenario_config) rep.orchestrator.step() def generate_distraction_scenarios(self, num_samples: int = 2000): """ 生成分心场景数据 分心类型: - 手机使用(手持、打字) - 视线偏离(看侧面、看后视镜) - 操作设备(调空调、导航) """ for i in range(num_samples // 2): self._place_phone_in_hand() self._set_hand_gesture('typing') rep.orchestrator.step() for i in range(num_samples // 2): self._set_gaze_direction( yaw=np.random.uniform(-60, 60), pitch=np.random.uniform(-30, 30) ) rep.orchestrator.step() def generate_occlusion_scenarios(self, num_samples: int = 500): """ 生成遮挡场景 遮挡类型: - 墨镜 - 口罩 - 手遮挡 - 头发遮挡 """ self._attach_sunglasses() for i in range(num_samples // 4): rep.orchestrator.step() self._attach_mask() for i in range(num_samples // 4): rep.orchestrator.step() def setup_annotation(self): """设置自动标注""" bbox_2d = rep.annotators.bbox_2d() bbox_3d = rep.annotators.bbox_3d() segmentation = rep.annotators.segmentation() keypoints = rep.annotators.keypoints_2d( skeleton_definition='coco_wholebody' ) rep.WriterRegistry.register_annotator(bbox_2d) rep.WriterRegistry.register_annotator(bbox_3d) rep.WriterRegistry.register_annotator(segmentation) rep.WriterRegistry.register_annotator(keypoints) def export_dataset(self, output_dir: str, format: str = 'coco'): """ 导出数据集 Args: output_dir: 输出目录 format: 数据集格式(coco, yolo, kitti) """ writer = rep.writers.get_writer(format) writer.initialize(output_dir=output_dir) rep.orchestrator.run(writer=writer) print(f"数据集已导出到: {output_dir}")
if __name__ == "__main__": generator = CabinSyntheticDataGenerator('cabin.usd') generator.setup_randomization() generator.generate_fatigue_scenarios(1000) generator.generate_distraction_scenarios(2000) generator.generate_occlusion_scenarios(500) generator.export_dataset('/data/dms_synthetic')
|