refactor(verification): batch pipeline inference with progress tracking

- 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
This commit is contained in:
2026-04-11 17:09:11 +08:00
parent 79b49f122a
commit 8a56c9649d
2 changed files with 91 additions and 35 deletions

View File

@@ -1,6 +1,6 @@
remote_root := "ial-gpu-workstation-1:/home/ial-pangyg/docker-workspace/projects/mini-nav" remote_root := "ial-gpu-workstation-1:/home/ial-pangyg/docker-workspace/projects/mini-nav"
rsync_flags := "-avLh --progress --stats --itemize-changes" 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: upload:
if command -v cygpath >/dev/null 2>&1 || test -n "${MSYSTEM:-}" || test -n "${CYGWIN:-}"; then \ if command -v cygpath >/dev/null 2>&1 || test -n "${MSYSTEM:-}" || test -n "${CYGWIN:-}"; then \

View File

@@ -69,9 +69,9 @@ def project_imports():
def habitat_setup(HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np): def habitat_setup(HabitatSimulatorConfig, RoomNode, create_habitat_simulator, np):
"""Initialize Habitat simulator and sample room nodes.""" """Initialize Habitat simulator and sample room nodes."""
_scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb" _scene_path = "data/scene_datasets/habitat-test-scenes/skokloster-castle.glb"
_image_size = 512 _image_size = 768
_num_rooms = 4 _num_rooms = 5
views_per_room = 6 views_per_room = 12
meters_per_pixel = 0.05 meters_per_pixel = 0.05
sim, agent = create_habitat_simulator( sim, agent = create_habitat_simulator(
@@ -166,12 +166,14 @@ def build_scene_graph(
ObjectNode, ObjectNode,
SimpleSceneGraph, SimpleSceneGraph,
cfg_manager, cfg_manager,
mo,
np, np,
pipeline, pipeline,
room_nodes, room_nodes,
room_view_dataset, room_view_dataset,
save_object_image, save_object_image,
save_room_view, save_room_view,
torch,
): ):
scene_graph = SimpleSceneGraph( scene_graph = SimpleSceneGraph(
rooms={_room.room_id: _room for _room in room_nodes}, rooms={_room.room_id: _room for _room in room_nodes},
@@ -199,23 +201,65 @@ def build_scene_graph(
"a door", "a door",
"a plant", "a plant",
] ]
_output = pipeline.process_batch( inference_batch_size = 4
_images, _text_labels, batch_size=32, return_debug_details=True 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
) )
_cropped_images = _output.cropped_images
hash_tensor = _output.hash_bits
from collections import Counter 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}") print(f"Fallback breakdown: {_reasons}")
# Save original room views. # 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) save_room_view(output_dir, _room_id, _view_idx, _image)
# Prefix sum: map flat crop index to (input_image_idx, mask_idx). # 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), ( assert sum(_num_selected) == len(_cropped_images), (
f"Sum of num_selected ({sum(_num_selected)}) != cropped_images count ({len(_cropped_images)})" f"Sum of num_selected ({sum(_num_selected)}) != cropped_images count ({len(_cropped_images)})"
) )
@@ -224,9 +268,23 @@ def build_scene_graph(
_prefix_sums.append(_prefix_sums[-1] + _n) _prefix_sums.append(_prefix_sums[-1] + _n)
_obj_counter = 0 _obj_counter = 0
_total_crops = len(_cropped_images)
object_tasks = []
for _img_idx, _n_crops in enumerate(_num_selected): for _img_idx, _n_crops in enumerate(_num_selected):
_room_id, _view_idx = _metadata[_img_idx] _room_id, _view_idx = _metadata[_img_idx]
for _mask_idx in range(_n_crops): for _mask_idx in range(_n_crops):
object_tasks.append((_img_idx, _room_id, _view_idx, _mask_idx, _n_crops))
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 _crop_flat_idx = _prefix_sums[_img_idx] + _mask_idx
_cropped = _cropped_images[_crop_flat_idx] _cropped = _cropped_images[_crop_flat_idx]
_hash_bits = hash_tensor[_crop_flat_idx] _hash_bits = hash_tensor[_crop_flat_idx]
@@ -238,9 +296,7 @@ def build_scene_graph(
_hash_bytes = np.packbits(_bits_binary).tobytes() _hash_bytes = np.packbits(_bits_binary).tobytes()
object_images[_obj_id] = _cropped object_images[_obj_id] = _cropped
save_object_image( save_object_image(output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped)
output_dir, _room_id, _obj_id, _view_idx, _mask_idx, _cropped
)
scene_graph.objects[_obj_id] = ObjectNode( scene_graph.objects[_obj_id] = ObjectNode(
obj_id=_obj_id, obj_id=_obj_id,
@@ -254,12 +310,12 @@ def build_scene_graph(
_obj_counter += 1 _obj_counter += 1
_fallback_count = sum( _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"Created {_obj_counter} objects")
print(f"Saved cropped images to: {output_dir}") 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 return hash_tensor, object_images, output_dir, scene_graph