From 8a56c9649d876f86e2a57ec529a26ce51232bd29 Mon Sep 17 00:00:00 2001 From: SikongJueluo Date: Sat, 11 Apr 2026 17:09:11 +0800 Subject: [PATCH] refactor(verification): batch pipeline inference with progress tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Increase verification params: image_size 512→768, rooms 4→5, views 6→12 - Refactor single-batch inference to chunked batch processing with mo.progress_bar - Extract debug_meta and hash_batches from output for clearer variable flow - Add progress bars to room-view snapshot saving and scene graph building - Add .ruff_cache/, .pytest_cache/, .sisyphus/ to .justfile upload excludes --- .justfile | 2 +- notebooks/verification.py | 124 +++++++++++++++++++++++++++----------- 2 files changed, 91 insertions(+), 35 deletions(-) diff --git a/.justfile b/.justfile index 3df223b..1304064 100644 --- a/.justfile +++ b/.justfile @@ -1,6 +1,6 @@ remote_root := "ial-gpu-workstation-1:/home/ial-pangyg/docker-workspace/projects/mini-nav" rsync_flags := "-avLh --progress --stats --itemize-changes" -upload_excludes := "--exclude=.jj/ --exclude=.git/ --exclude=.devenv/ --exclude=.direnv/ --exclude=deps/ --exclude=outputs/ --exclude=data/versioned_data/ --exclude=datasets/" +upload_excludes := "--exclude=.jj/ --exclude=.git/ --exclude=.devenv/ --exclude=.direnv/ --exclude=deps/ --exclude=outputs/ --exclude=data/versioned_data/ --exclude=datasets/ --exclude=.ruff_cache/ --exclude=.pytest_cache/ --exclude=.sisyphus/" upload: if command -v cygpath >/dev/null 2>&1 || test -n "${MSYSTEM:-}" || test -n "${CYGWIN:-}"; then \ diff --git a/notebooks/verification.py b/notebooks/verification.py index 2b2e81d..516b6ae 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -69,9 +69,9 @@ def project_imports(): def habitat_setup(HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np): """Initialize Habitat simulator and sample room nodes.""" _scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb" - _image_size = 512 - _num_rooms = 4 - views_per_room = 6 + _image_size = 768 + _num_rooms = 5 + views_per_room = 12 meters_per_pixel = 0.05 sim, agent = create_habitat_simulator( @@ -166,12 +166,14 @@ def build_scene_graph( ObjectNode, SimpleSceneGraph, cfg_manager, + mo, np, pipeline, room_nodes, room_view_dataset, save_object_image, save_room_view, + torch, ): scene_graph = SimpleSceneGraph( rooms={_room.room_id: _room for _room in room_nodes}, @@ -199,23 +201,65 @@ def build_scene_graph( "a door", "a plant", ] - _output = pipeline.process_batch( - _images, _text_labels, batch_size=32, return_debug_details=True - ) - _cropped_images = _output.cropped_images - hash_tensor = _output.hash_bits + inference_batch_size = 4 + image_batches = [ + _images[index : index + inference_batch_size] + for index in range(0, len(_images), inference_batch_size) + ] + + _cropped_images = [] + debug_meta = [] + hash_batches = [] + for _batch_images in mo.status.progress_bar( + image_batches, + title="Running pipeline inference on room views", + subtitle=f"Batch size {inference_batch_size} with ETA", + completion_title="Pipeline inference finished", + completion_subtitle=( + f"Processed {len(_images)} room views in {len(image_batches)} batches" + ), + show_eta=True, + show_rate=True, + remove_on_exit=False, + ): + _batch_output = pipeline.process_batch( + _batch_images, + _text_labels, + batch_size=inference_batch_size, + return_debug_details=True, + ) + _cropped_images.extend(_batch_output.cropped_images) + debug_meta.extend(_batch_output.debug_meta) + if _batch_output.hash_bits.numel() > 0: + hash_batches.append(_batch_output.hash_bits) + + if hash_batches: + hash_tensor = torch.cat(hash_batches, dim=0) + else: + hash_tensor = torch.empty( + (0, pipeline.hash_bits), dtype=torch.int32, device=pipeline.device + ) from collections import Counter - _reasons = Counter(m["fallback_reason"] or "ok" for m in _output.debug_meta) + _reasons = Counter(m["fallback_reason"] or "ok" for m in debug_meta) print(f"Fallback breakdown: {_reasons}") # Save original room views. - for _room_id, _view_idx, _image in room_view_dataset: + for _room_id, _view_idx, _image in mo.status.progress_bar( + room_view_dataset, + title="Saving room-view snapshots", + subtitle="Writing original room images to disk", + completion_title="Room-view snapshots saved", + completion_subtitle=f"Saved {len(room_view_dataset)} room views", + show_eta=True, + show_rate=True, + remove_on_exit=False, + ): save_room_view(output_dir, _room_id, _view_idx, _image) # Prefix sum: map flat crop index to (input_image_idx, mask_idx). - _num_selected = [_m["num_selected"] for _m in _output.debug_meta] + _num_selected = [_m["num_selected"] for _m in debug_meta] assert sum(_num_selected) == len(_cropped_images), ( f"Sum of num_selected ({sum(_num_selected)}) != cropped_images count ({len(_cropped_images)})" ) @@ -224,42 +268,54 @@ def build_scene_graph( _prefix_sums.append(_prefix_sums[-1] + _n) _obj_counter = 0 + _total_crops = len(_cropped_images) + object_tasks = [] for _img_idx, _n_crops in enumerate(_num_selected): _room_id, _view_idx = _metadata[_img_idx] for _mask_idx in range(_n_crops): - _crop_flat_idx = _prefix_sums[_img_idx] + _mask_idx - _cropped = _cropped_images[_crop_flat_idx] - _hash_bits = hash_tensor[_crop_flat_idx] + object_tasks.append((_img_idx, _room_id, _view_idx, _mask_idx, _n_crops)) - _obj_id = f"{_room_id}_v{_view_idx:03d}_m{_mask_idx:02d}" + for _img_idx, _room_id, _view_idx, _mask_idx, _n_crops in mo.status.progress_bar( + object_tasks, + title="Building scene graph objects", + subtitle="Preparing cropped objects and hashes with ETA", + completion_title="Scene graph build complete", + completion_subtitle=f"Created {_total_crops} cropped object entries", + show_eta=True, + show_rate=True, + remove_on_exit=False, + ): + _crop_flat_idx = _prefix_sums[_img_idx] + _mask_idx + _cropped = _cropped_images[_crop_flat_idx] + _hash_bits = hash_tensor[_crop_flat_idx] - _bits_array = _hash_bits.detach().cpu().numpy().reshape(-1) - _bits_binary = (_bits_array > 0).astype(np.uint8) - _hash_bytes = np.packbits(_bits_binary).tobytes() + _obj_id = f"{_room_id}_v{_view_idx:03d}_m{_mask_idx:02d}" - object_images[_obj_id] = _cropped - save_object_image( - output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped - ) + _bits_array = _hash_bits.detach().cpu().numpy().reshape(-1) + _bits_binary = (_bits_array > 0).astype(np.uint8) + _hash_bytes = np.packbits(_bits_binary).tobytes() - scene_graph.objects[_obj_id] = ObjectNode( - obj_id=_obj_id, - room_id=_room_id, - position=np.array([0.0, 0.0, 0.0], dtype=np.float32), - visual_hash=_hash_bytes, - semantic_hash=_hash_bytes, - hit_count=1, - last_seen_frame=_view_idx, - ) - _obj_counter += 1 + object_images[_obj_id] = _cropped + save_object_image(output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped) + + scene_graph.objects[_obj_id] = ObjectNode( + obj_id=_obj_id, + room_id=_room_id, + position=np.array([0.0, 0.0, 0.0], dtype=np.float32), + visual_hash=_hash_bytes, + semantic_hash=_hash_bytes, + hit_count=1, + last_seen_frame=_view_idx, + ) + _obj_counter += 1 _fallback_count = sum( - 1 for _meta in _output.debug_meta if _meta["fallback_reason"] is not None + 1 for _meta in debug_meta if _meta["fallback_reason"] is not None ) print(f"Created {_obj_counter} objects") print(f"Saved cropped images to: {output_dir}") - print(f"Fallback frames: {_fallback_count}/{len(_output.debug_meta)}") + print(f"Fallback frames: {_fallback_count}/{len(debug_meta)}") return hash_tensor, object_images, output_dir, scene_graph