(rng, out_dir: Path,
kubasic_source, gso_source, hdri_source,
objects_split="train", backgrounds_split="train",
frame_start=1, frame_end=40, frame_rate=12,
resolution=256, step_rate=240, spp=DEFAULT_SPP)
| 160 | subprocess.run(cmd, check=True) |
| 161 | |
| 162 | def generate_one(rng, out_dir: Path, |
| 163 | kubasic_source, gso_source, hdri_source, |
| 164 | objects_split="train", backgrounds_split="train", |
| 165 | frame_start=1, frame_end=40, frame_rate=12, |
| 166 | resolution=256, step_rate=240, spp=DEFAULT_SPP): |
| 167 | |
| 168 | print(f"\n[DEBUG] ===== GENERATING SCENE: {out_dir.name} =====") |
| 169 | |
| 170 | # RANDOMIZE NUMBER OF OBJECTS (2-5) |
| 171 | num_objects = rng.choice([3, 3, 4, 4, 5, 6, 6, 7, 7]) # Weighted toward 3-4 objects |
| 172 | num_to_remove = 2 if num_objects >= 4 else 1 |
| 173 | |
| 174 | print(f"[DEBUG] Scene will have {num_objects} objects, removing {num_to_remove}") |
| 175 | |
| 176 | tmp_scratch = tempfile.mkdtemp(prefix="kb_scratch_") |
| 177 | FLAGS_like = NS( |
| 178 | frame_start=frame_start, |
| 179 | frame_end=frame_end, |
| 180 | frame_rate=frame_rate, |
| 181 | step_rate=step_rate, |
| 182 | resolution=resolution, |
| 183 | job_dir=str(out_dir.parent), |
| 184 | scratch_dir=tmp_scratch, |
| 185 | seed=12, |
| 186 | logging_level="WARNING", |
| 187 | ) |
| 188 | |
| 189 | tmpA = tmpB = tmpC = tmpMask = None |
| 190 | try: |
| 191 | cleanup_scene() |
| 192 | |
| 193 | # ========== PASS A: All objects (full physics, all visible) ========== |
| 194 | print(f"[DEBUG] Creating Pass A with {num_objects} objects...") |
| 195 | |
| 196 | scene, _, output_dir, scratch_dir = kb.setup(FLAGS_like) |
| 197 | renderer = Blender(scene, scratch_dir, samples_per_pixel=int(spp)) |
| 198 | simulator = PyBullet(scene, scratch_dir) |
| 199 | |
| 200 | # Background (bright only) |
| 201 | train_bg, test_bg = hdri_source.get_test_split(fraction=0.1) |
| 202 | bg_candidates = train_bg if backgrounds_split=="train" else test_bg |
| 203 | bright_backgrounds = [bg for bg in bg_candidates if not any(dark in bg for dark in DARK_BACKGROUNDS)] |
| 204 | if not bright_backgrounds: |
| 205 | bright_backgrounds = bg_candidates |
| 206 | |
| 207 | hdri_id = rng.choice(bright_backgrounds) |
| 208 | bg_tex = hdri_source.create(asset_id=hdri_id) |
| 209 | scene.metadata["background"] = hdri_id |
| 210 | renderer._set_ambient_light_hdri(bg_tex.filename) |
| 211 | print(f"[DEBUG] Background: {hdri_id}") |
| 212 | |
| 213 | # Dome (floor) |
| 214 | dome = kubasic_source.create(asset_id="dome", name="dome", |
| 215 | friction=0.5, restitution=0.2, |
| 216 | static=True, background=True) |
| 217 | scene += dome |
| 218 | dome_b = dome.linked_objects[renderer] |
| 219 | dome_b.data.materials[0].node_tree.nodes["Image Texture"].image = bpy.data.images.load(bg_tex.filename) |
no test coverage detected