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
| import omni.replicator.core as rep from omni.isaac.sim import SimulationApp
class DMSDataGenerator: """ DMS合成数据生成器 """ def __init__(self, headless=True): self.sim = SimulationApp({"headless": headless}) self.config = { "resolution": (1920, 1080), "fps": 30, "output_dir": "/data/dms_synthetic" } def setup_cabin_environment(self): """ 配置座舱环境 """ import omni from omni.isaac.core.utils.stage import add_reference_to_stage cabin_asset = "omniverse://localhost/NVIDIA/Assets/Vehicles/Cabin.usd" add_reference_to_stage(cabin_asset, "/World/Cabin") self.setup_lighting() self.setup_camera() def setup_camera(self): """ 配置DMS相机(仪表板位置) """ from omni.isaac.sensor import Camera self.camera = Camera( prim_path="/World/Cabin/DMS_Camera", position=(0.2, 0.0, 1.2), orientation=(0, 0, 0, 1), frequency=self.config["fps"], resolution=self.config["resolution"] ) def setup_driver_model(self): """ 配置驾驶员模型(Metahuman) """ driver_params = { "gender": ["male", "female"], "age": [20, 30, 40, 50, 60], "ethnicity": ["asian", "caucasian", "african"], "glasses": [True, False], "facial_hair": [True, False] } return driver_params def randomize_scene(self): """ 场景随机化 """ time_of_day = rep.random.uniform(6, 20) light_intensity = rep.random.uniform(0.5, 1.5) head_pose = { "yaw": rep.random.uniform(-30, 30), "pitch": rep.random.uniform(-20, 20), "roll": rep.random.uniform(-10, 10) } gaze_direction = rep.random.uniform_2d(-0.5, 0.5) expression = rep.random.choice(["neutral", "tired", "angry"]) return { "time_of_day": time_of_day, "light_intensity": light_intensity, "head_pose": head_pose, "gaze_direction": gaze_direction, "expression": expression } def generate_dataset(self, num_samples=10000): """ 生成DMS数据集 """ writer = rep.WriterRegistry.get("BasicWriter") writer.initialize( output_dir=self.config["output_dir"], rgb=True, bounding_boxes=True, segmentation=True, keypoints=True ) for i in range(num_samples): self.randomize_scene() self.sim.update() writer.write(frame=i) print(f"生成 {num_samples} 张图像到 {self.config['output_dir']}")
if __name__ == "__main__": generator = DMSDataGenerator(headless=True) generator.setup_cabin_environment() generator.generate_dataset(num_samples=10000)
|