feat(compressors): refactor pipeline with FramePacket dataclass and unified process_batch

- Add FramePacket dataclass to encapsulate per-image pipeline state
- Rename internal methods with underscore prefix convention
- Replace separate filter_batch/crop_batch with unified process_batch method
- Update notebook to use new HashPipeline API
This commit is contained in:
2026-04-04 19:55:36 +08:00
parent 94ed05a039
commit 3638ffdb8d
3 changed files with 450 additions and 243 deletions

View File

@@ -160,7 +160,6 @@ def collect_views(
@app.cell
def build_scene_graph(
Image,
ObjectNode,
SimpleSceneGraph,
cfg_manager,
@@ -168,9 +167,7 @@ def build_scene_graph(
pipeline,
room_nodes,
room_view_dataset,
torch,
):
"""Build scene graph using step-by-step pipeline to capture cropped images."""
scene_graph = SimpleSceneGraph(
rooms={_room.room_id: _room for _room in room_nodes},
objects={},
@@ -185,36 +182,10 @@ def build_scene_graph(
_images = [item[2] for item in room_view_dataset]
_metadata = [(item[0], item[1]) for item in room_view_dataset]
# Step 1: Detect objects.
_text_labels = ["object"]
_detections = pipeline.detect_batch(_images, _text_labels)
# Step 2: Segment with SAM.
_bboxes_per_image = [[_d["bbox"] for _d in _dets] for _dets in _detections]
_masks = pipeline.segment_batch(_images, _bboxes_per_image)
# Step 3: Filter masks.
_filtered = pipeline.filter_batch(_images, _masks)
# Step 4: Crop images.
_cropped_images = pipeline.crop_batch(_filtered, _masks, _detections)
# Step 5: Extract DINO features and compress to hash.
_batch_size = 32
_all_bits = []
for _i in range(0, len(_cropped_images), _batch_size):
_batch = _cropped_images[_i : _i + _batch_size]
_tokens = pipeline.extract_dino_batch(_batch)
_bits = pipeline.compress_batch(_tokens)
_all_bits.append(_bits)
hash_tensor = (
torch.cat(_all_bits, dim=0)
if _all_bits
else torch.empty(
(0, pipeline.hash_bits), dtype=torch.int32, device=pipeline.device
)
)
_output = pipeline.process_batch(_images, _text_labels, batch_size=32)
_cropped_images = _output.cropped_images
hash_tensor = _output.hash_bits
# Step 6: Create ObjectNodes and save cropped images.
for _idx, (_cropped, _hash_bits) in enumerate(zip(_cropped_images, hash_tensor)):
@@ -238,8 +209,13 @@ def build_scene_graph(
last_seen_frame=_view_idx,
)
_fallback_count = sum(
1 for _meta in _output.debug_meta if _meta["fallback_reason"] is not None
)
print(f"Created {len(scene_graph.objects)} objects")
print(f"Saved cropped images to: {output_dir}")
print(f"Fallback frames: {_fallback_count}/{len(_output.debug_meta)}")
return hash_tensor, object_images, output_dir, scene_graph
@@ -302,19 +278,12 @@ def query_matching(
if file_upload.value:
_query_image = Image.open(io.BytesIO(file_upload.contents())).convert("RGB")
# Step-by-step processing to get cropped query image.
_text_labels = ["object"]
_detections = pipeline.detect_batch([_query_image], _text_labels)
_bboxes = [[_d["bbox"] for _d in _dets] for _dets in _detections]
_masks = pipeline.segment_batch([_query_image], _bboxes)
_filtered = pipeline.filter_batch([_query_image], _masks)
_cropped = pipeline.crop_batch(_filtered, _masks, _detections)
_tokens = pipeline.extract_dino_batch(_cropped)
_query_bits = pipeline.compress_batch(_tokens)
_output = pipeline.process_batch([_query_image], _text_labels, batch_size=1)
_query_bits = _output.hash_bits
if _query_bits.numel() > 0:
query_cropped = _cropped[0]
query_cropped = _output.cropped_images[0]
_query_tensor = _query_bits[0].int()
_obj_ids = list(scene_graph.objects.keys())