mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
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:
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user