()
| 607 | gc.collect() |
| 608 | |
| 609 | def main(): |
| 610 | parser = kb.ArgumentParser() |
| 611 | parser.add_argument("--objects_split", choices=["train","test"], default="train") |
| 612 | parser.add_argument("--backgrounds_split", choices=["train","test"], default="train") |
| 613 | parser.add_argument("--kubasic_assets", type=str, default="gs://kubric-public/assets/KuBasic/KuBasic.json") |
| 614 | parser.add_argument("--gso_assets", type=str, default="gs://kubric-public/assets/GSO/GSO.json") |
| 615 | parser.add_argument("--hdri_assets", type=str, default="gs://kubric-public/assets/HDRI_haven/HDRI_haven.json") |
| 616 | parser.set_defaults(frame_start=1, frame_end=60, frame_rate=12, resolution=384, step_rate=240) |
| 617 | parser.add_argument("--out_prefix", type=str, default="variable_obj_3") |
| 618 | parser.add_argument("--num_pairs", type=int, default=200) |
| 619 | parser.add_argument("--start_index", type=int, default=0) |
| 620 | parser.add_argument("--fast", action="store_true") |
| 621 | |
| 622 | FLAGS = parser.parse_args() |
| 623 | |
| 624 | if FLAGS.fast: |
| 625 | FLAGS.frame_end = 24 |
| 626 | FLAGS.frame_rate = 8 |
| 627 | spp = 16 |
| 628 | else: |
| 629 | spp = DEFAULT_SPP |
| 630 | |
| 631 | # Load assets |
| 632 | print("[MAIN] Loading assets...") |
| 633 | kubasic_source = kb.AssetSource.from_manifest(FLAGS.kubasic_assets) |
| 634 | gso_source = kb.AssetSource.from_manifest(FLAGS.gso_assets) |
| 635 | hdri_source = kb.AssetSource.from_manifest(FLAGS.hdri_assets) |
| 636 | |
| 637 | _, rng, output_dir, _ = kb.setup(FLAGS) |
| 638 | root = output_dir / FLAGS.out_prefix |
| 639 | root.mkdir(parents=True, exist_ok=True) |
| 640 | |
| 641 | print(f"[MAIN] Generating {FLAGS.num_pairs} video pairs with variable objects...") |
| 642 | successful = 0 |
| 643 | failed = 0 |
| 644 | |
| 645 | for i in range(FLAGS.start_index, FLAGS.start_index + FLAGS.num_pairs): |
| 646 | out_dir = root / f"{i:05d}" |
| 647 | out_dir.mkdir(parents=True, exist_ok=True) |
| 648 | |
| 649 | local_rng = np.random.RandomState(i + 2001) |
| 650 | print(f"\n[BATCH {i:05d}/{FLAGS.start_index + FLAGS.num_pairs - 1:05d}] Starting...") |
| 651 | |
| 652 | try: |
| 653 | generate_one(local_rng, out_dir, |
| 654 | kubasic_source, gso_source, hdri_source, |
| 655 | objects_split=FLAGS.objects_split, |
| 656 | backgrounds_split=FLAGS.backgrounds_split, |
| 657 | frame_start=FLAGS.frame_start, |
| 658 | frame_end=FLAGS.frame_end, |
| 659 | frame_rate=FLAGS.frame_rate, |
| 660 | resolution=FLAGS.resolution, |
| 661 | step_rate=FLAGS.step_rate, |
| 662 | spp=spp) |
| 663 | print(f"[BATCH {i:05d}] ✓ Success") |
| 664 | successful += 1 |
| 665 | except Exception as e: |
| 666 | print(f"[BATCH {i:05d}] ✗ ERROR: {e}") |
no test coverage detected