diff --git a/mini-nav/compressors/pipeline.py b/mini-nav/compressors/pipeline.py index e1f71d0..6920c66 100644 --- a/mini-nav/compressors/pipeline.py +++ b/mini-nav/compressors/pipeline.py @@ -1,5 +1,6 @@ """OWLv2 + SAM + DINO + Hash compression pipeline.""" +from dataclasses import dataclass, field from typing import Any, Optional, Sequence import torch @@ -8,7 +9,12 @@ from PIL import Image from utils import get_device from utils.image import crop_image_by_bbox, extract_masked_region -from .filter import MaskScoringConfig, select_best_mask +from .filter import ( + MaskScoringConfig, + compute_mask_features, + score_mask, + should_reject_mask, +) from .model_loader import ( get_dino_dim, load_dino_model, @@ -24,14 +30,6 @@ from .proposal.core import DetectionResult def create_pipeline_from_config(config) -> "HashPipeline": - """Create HashPipeline from a config object. - - Args: - config: Configuration object with model settings. - - Returns: - Initialized HashPipeline. - """ return HashPipeline( owlv2_model=getattr( config.model, "owlv2_model", "google/owlv2-base-patch16-ensemble" @@ -46,18 +44,40 @@ def create_pipeline_from_config(config) -> "HashPipeline": ) +@dataclass +class FramePacket: + image: Image.Image + boxes_xyxy: list[list[float]] = field(default_factory=list) + scores: list[float] = field(default_factory=list) + labels: list[str] = field(default_factory=list) + masks: list[dict[str, Any]] = field(default_factory=list) + selected_idx: int | None = None + dropped_indices: list[int] = field(default_factory=list) + fallback_reason: str | None = None + filtered_image: Image.Image | None = None + cropped_image: Image.Image | None = None + + +@dataclass +class PipelineBatchOutput: + hash_bits: torch.Tensor + cropped_images: list[Image.Image] + debug_meta: list[dict[str, Any]] + + class HashPipeline: """Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression. Pipeline flow: Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) -> - Crop (OWLv2 box) -> DINO (features) -> Hash (binary codes) + Crop (matched OWLv2 box) -> DINO (features) -> Hash (binary codes) Example: pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512) images = [Image.open("path/to/image.jpg")] text_labels = ["object"] - hash_bits = pipeline.forward_batch(images, text_labels) # Returns [N, 512] + output = pipeline.process_batch(images, text_labels) + hash_bits = output.hash_bits # [N, 512] """ def __init__( @@ -73,27 +93,16 @@ class HashPipeline: ): super().__init__() - # Device for model placement. self.device = get_device() - - # OWLv2 detection settings. self.owlv2_processor, self.owlv2_model = load_owlv2_model( model_name=owlv2_model ) self.score_threshold = score_threshold self.postprocess_threshold = postprocess_threshold - - # Mask scoring config for filter step. self.mask_scoring_config = mask_scoring_config - - # SAM2 model for segmentation. self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model) - - # DINO model for feature extraction. self.dino_processor, self.dino = load_dino_model(model_name=dino_model) self.dino_dim = get_dino_dim(dino_model) - - # Hash compressor for binarizing DINO features. self.hash_compressor = load_hash_compressor( input_dim=self.dino_dim, hash_bits=hash_bits, @@ -102,23 +111,13 @@ class HashPipeline: @property def hash_bits(self) -> int: - """Number of bits in the hash code.""" return self.hash_compressor.hash_bits - def detect_batch( + def _detect_batch( self, images: Sequence[Image.Image], text_labels: list[str], ) -> list[list[DetectionResult]]: - """Detect objects in a batch of images using OWLv2. - - Args: - images: Sequence of PIL Images. - text_labels: Text labels for all images (same labels used for each image). - - Returns: - List of lists of DetectionResult dicts, one inner list per image. - """ image_list = list(images) if not image_list: return [] @@ -133,20 +132,11 @@ class HashPipeline: postprocess_threshold=self.postprocess_threshold, ) - def segment_batch( + def _segment_batch( self, images: Sequence[Image.Image], bboxes_per_image: list[list[list[float]]], ) -> list[list[dict[str, Any]]]: - """Segment objects in images using SAM2 with bounding box prompts. - - Args: - images: Sequence of PIL Images. - bboxes_per_image: Bounding boxes per image as [[[x1,y1,x2,y2], ...], ...]. - - Returns: - List of lists of mask dictionaries, one inner list per image. - """ image_list = list(images) if not image_list: return [] @@ -158,77 +148,169 @@ class HashPipeline: bboxes_per_image, ) - def filter_batch( + def _build_frame_packets(self, images: Sequence[Image.Image]) -> list[FramePacket]: + return [FramePacket(image=image) for image in images] + + def _attach_detections( self, - images: Sequence[Image.Image], - masks_per_image: list[list[dict[str, Any]]], - ) -> list[Image.Image]: - """Filter masks and extract best masked region for each image. - - Args: - images: Sequence of PIL Images. - masks_per_image: Masks per image from segment_batch. - - Returns: - List of PIL Images, one per input image (original if no valid masks). - """ - image_list = list(images) - if not image_list: + packets: list[FramePacket], + text_labels: list[str], + ) -> list[list[DetectionResult]]: + if not packets: return [] - filtered_images: list[Image.Image] = [] - for index, image in enumerate(image_list): - masks = masks_per_image[index] if index < len(masks_per_image) else [] - if not masks: - filtered_images.append(image) - continue - - best_mask = select_best_mask( - masks, - image_shape=(image.height, image.width), - config=self.mask_scoring_config, - ) - if best_mask is None: - filtered_images.append(image) - continue - - filtered_images.append(extract_masked_region(image, best_mask["segment"])) - - return filtered_images - - def crop_batch( - self, - images: Sequence[Image.Image], - masks_per_image: list[list[dict[str, Any]]], - detections_per_image: list[list[DetectionResult]], - ) -> list[Image.Image]: - """Crop filtered images using OWLv2 detection boxes. - - Args: - images: Sequence of PIL Images after filter_batch. - masks_per_image: Masks per image from segment_batch. - detections_per_image: Detection results per image from detect_batch. - - Returns: - List of cropped PIL Images. Returns original image when no detection exists. - """ - image_list = list(images) - if not image_list: - return [] - - cropped_images: list[Image.Image] = [] - for index, image in enumerate(image_list): + image_list = [packet.image for packet in packets] + detections_per_image = self._detect_batch(image_list, text_labels) + for index, packet in enumerate(packets): detections = ( detections_per_image[index] if index < len(detections_per_image) else [] ) - if detections: - best_detection = max(detections, key=lambda d: d["score"]) - cropped_images.append(crop_image_by_bbox(image, best_detection["bbox"])) + packet.boxes_xyxy = [list(det["bbox"]) for det in detections] + packet.scores = [float(det["score"]) for det in detections] + packet.labels = [str(det["label"]) for det in detections] + if not packet.boxes_xyxy: + packet.fallback_reason = "no_detection" + return detections_per_image + + def _attach_masks(self, packets: list[FramePacket]) -> None: + if not packets: + return + + image_list = [packet.image for packet in packets] + boxes_per_image = [packet.boxes_xyxy for packet in packets] + masks_per_image = self._segment_batch(image_list, boxes_per_image) + + for index, packet in enumerate(packets): + packet.masks = ( + masks_per_image[index] if index < len(masks_per_image) else [] + ) + if ( + packet.boxes_xyxy + and not packet.masks + and packet.fallback_reason is None + ): + packet.fallback_reason = "no_mask" + + def _select_candidates(self, packets: list[FramePacket]) -> None: + config = self.mask_scoring_config or MaskScoringConfig() + + for packet in packets: + if not packet.masks: + packet.selected_idx = None + packet.dropped_indices = [] + if packet.fallback_reason is None: + packet.fallback_reason = "no_mask" continue - cropped_images.append(image) + kept: list[tuple[float, int, int]] = [] + dropped: list[int] = [] + for index, mask in enumerate(packet.masks): + features = compute_mask_features( + mask, image_shape=(packet.image.height, packet.image.width) + ) + if should_reject_mask(features, config): + dropped.append(index) + continue - return cropped_images + mask_score = score_mask( + mask, + image_shape=(packet.image.height, packet.image.width), + config=config, + ) + area = int(mask.get("area", 0)) + kept.append((float(mask_score), area, index)) + + if kept: + kept.sort(reverse=True) + packet.selected_idx = kept[0][2] + packet.dropped_indices = dropped + continue + + fallback_index = max( + range(len(packet.masks)), + key=lambda idx: int(packet.masks[idx].get("area", 0)), + ) + packet.selected_idx = fallback_index + packet.dropped_indices = [ + index for index in range(len(packet.masks)) if index != fallback_index + ] + if packet.fallback_reason is None: + packet.fallback_reason = "all_masks_rejected_fallback_area" + + def _render_filtered_images(self, packets: list[FramePacket]) -> None: + for packet in packets: + if packet.selected_idx is None: + packet.filtered_image = packet.image + continue + + if packet.selected_idx >= len(packet.masks): + packet.filtered_image = packet.image + packet.fallback_reason = ( + packet.fallback_reason or "selected_index_out_of_mask_range" + ) + packet.selected_idx = None + continue + + selected_mask = packet.masks[packet.selected_idx] + packet.filtered_image = extract_masked_region( + packet.image, selected_mask["segment"] + ) + + def _render_cropped_images( + self, + packets: list[FramePacket], + detections_per_image: list[list[DetectionResult]], + ) -> None: + for index, packet in enumerate(packets): + base_image = ( + packet.filtered_image if packet.filtered_image else packet.image + ) + detections = ( + detections_per_image[index] if index < len(detections_per_image) else [] + ) + + if packet.selected_idx is None: + packet.cropped_image = base_image + if packet.fallback_reason is None: + packet.fallback_reason = "no_selected_candidate" + continue + + if packet.selected_idx >= len(detections): + packet.cropped_image = base_image + packet.fallback_reason = ( + packet.fallback_reason or "selected_index_out_of_detection_range" + ) + continue + + selected_detection = detections[packet.selected_idx] + cropped = crop_image_by_bbox(base_image, selected_detection["bbox"]) + packet.cropped_image = cropped + if cropped.size == base_image.size: + packet.fallback_reason = ( + packet.fallback_reason or "invalid_or_full_bbox" + ) + + def _build_debug_meta( + self, + packets: list[FramePacket], + return_debug_details: bool, + ) -> list[dict[str, Any]]: + debug_meta: list[dict[str, Any]] = [] + for packet in packets: + item: dict[str, Any] = { + "selected_idx": packet.selected_idx, + "dropped_indices": list(packet.dropped_indices), + "fallback_reason": packet.fallback_reason, + "num_boxes": len(packet.boxes_xyxy), + "num_masks": len(packet.masks), + } + if return_debug_details: + item["boxes_xyxy"] = [list(box) for box in packet.boxes_xyxy] + item["scores"] = [float(score) for score in packet.scores] + item["labels"] = [str(label) for label in packet.labels] + item["masks"] = packet.masks + debug_meta.append(item) + return debug_meta def extract_dino_batch(self, images: Sequence[Image.Image]) -> torch.Tensor: """Extract DINO tokens from a batch of images. @@ -265,6 +347,65 @@ class HashPipeline: _, _, bits = self.hash_compressor(tokens) return bits + def process_batch( + self, + images: Sequence[Image.Image], + text_labels: list[str], + batch_size: int = 32, + return_debug_details: bool = False, + ) -> PipelineBatchOutput: + """Run full pipeline and return cropped images + hashes + debug metadata. + + Args: + images: Sequence of PIL Images. + text_labels: Text labels for detection (same for all images). + batch_size: Batch size for DINO feature extraction. + return_debug_details: Include boxes/scores/labels/masks in debug output. + + Returns: + PipelineBatchOutput with final cropped images, binary hash bits, + and per-image debug metadata. + """ + if batch_size <= 0: + raise ValueError("batch_size must be greater than 0") + + image_list = list(images) + if not image_list: + return PipelineBatchOutput( + hash_bits=torch.empty( + (0, self.hash_bits), dtype=torch.int32, device=self.device + ), + cropped_images=[], + debug_meta=[], + ) + + packets = self._build_frame_packets(image_list) + detections_per_image = self._attach_detections(packets, text_labels) + self._attach_masks(packets) + self._select_candidates(packets) + self._render_filtered_images(packets) + self._render_cropped_images(packets, detections_per_image) + + cropped_images = [ + packet.cropped_image if packet.cropped_image is not None else packet.image + for packet in packets + ] + + all_bits: list[torch.Tensor] = [] + for index in range(0, len(cropped_images), batch_size): + sub_batch = cropped_images[index : index + batch_size] + tokens = self.extract_dino_batch(sub_batch) + bits = self.compress_batch(tokens) + all_bits.append(bits) + + hash_bits = torch.cat(all_bits, dim=0) + debug_meta = self._build_debug_meta(packets, return_debug_details) + return PipelineBatchOutput( + hash_bits=hash_bits, + cropped_images=cropped_images, + debug_meta=debug_meta, + ) + def forward_batch( self, images: Sequence[Image.Image], @@ -281,29 +422,11 @@ class HashPipeline: Returns: Binary hash codes of shape [N, hash_bits] as int32. """ - if batch_size <= 0: - raise ValueError("batch_size must be greater than 0") - - image_list = list(images) - if not image_list: - return torch.empty( - (0, self.hash_bits), dtype=torch.int32, device=self.device - ) - - detections = self.detect_batch(image_list, text_labels) - bboxes = [[d["bbox"] for d in dets] for dets in detections] - masks = self.segment_batch(image_list, bboxes) - processed = self.filter_batch(image_list, masks) - processed = self.crop_batch(processed, masks, detections) - - all_bits: list[torch.Tensor] = [] - for i in range(0, len(processed), batch_size): - sub_batch = processed[i : i + batch_size] - tokens = self.extract_dino_batch(sub_batch) - bits = self.compress_batch(tokens) - all_bits.append(bits) - - return torch.cat(all_bits, dim=0) + return self.process_batch( + images=images, + text_labels=text_labels, + batch_size=batch_size, + ).hash_bits def extract_features_dataset( self, @@ -330,15 +453,15 @@ class HashPipeline: (0, self.dino_dim), dtype=torch.float32, device=self.device ) - detections = self.detect_batch(image_list, text_labels) - bboxes = [[d["bbox"] for d in dets] for dets in detections] - masks = self.segment_batch(image_list, bboxes) - processed = self.filter_batch(image_list, masks) - processed = self.crop_batch(processed, masks, detections) + processed = self.process_batch( + images=image_list, + text_labels=text_labels, + batch_size=batch_size, + ).cropped_images all_features: list[torch.Tensor] = [] - for i in range(0, len(processed), batch_size): - sub_batch = processed[i : i + batch_size] + for index in range(0, len(processed), batch_size): + sub_batch = processed[index : index + batch_size] tokens = self.extract_dino_batch(sub_batch) features = tokens.mean(dim=1) all_features.append(F.normalize(features, dim=-1)) diff --git a/notebooks/proposal_segament.py b/notebooks/proposal_segament.py index 188c594..b5f8dcb 100644 --- a/notebooks/proposal_segament.py +++ b/notebooks/proposal_segament.py @@ -76,8 +76,7 @@ def _(agent, room_nodes, sim, views_per_room): @app.cell def _(ImageDraw, ImageFont, all_room_views, mo): - from compressors.model_loader import load_owlv2_model - from compressors.proposal.core import detect_objects_batch + from compressors import HashPipeline from utils.common import get_device from utils.image import numpy_to_pil @@ -96,22 +95,25 @@ def _(ImageDraw, ImageFont, all_room_views, mo): "a window", ] - owl_processor, owl_model = load_owlv2_model( - model_name="google/owlv2-base-patch16-ensemble" + pipeline = HashPipeline( + owlv2_model="google/owlv2-base-patch16-ensemble", + score_threshold=score_threshold, + postprocess_threshold=postprocess_threshold, ) image = numpy_to_pil(all_room_views["room_00"][1]) - detection_batch = detect_objects_batch( - model=owl_model, - processor=owl_processor, + output = pipeline.process_batch( images=[image], - text_labels_per_image=[text_labels], - score_threshold=score_threshold, - postprocess_threshold=postprocess_threshold, + text_labels=text_labels, + batch_size=1, + return_debug_details=True, ) - detections = detection_batch[0] if detection_batch else [] - filtered_items = [(det["bbox"], det["score"], det["label"]) for det in detections] + meta = output.debug_meta[0] if output.debug_meta else {} + boxes = meta.get("boxes_xyxy", []) + scores = meta.get("scores", []) + labels = meta.get("labels", []) + filtered_items = list(zip(boxes, scores, labels, strict=False)) _vis_image = image.copy() _draw = ImageDraw.Draw(_vis_image) @@ -162,32 +164,56 @@ def _(ImageDraw, ImageFont, all_room_views, mo): mo.md(detection_text), ] ) - return device, filtered_items, image + return device, filtered_items, image, meta @app.cell -def _(Image, ImageDraw, device, filtered_items, image, mo, np): - from compressors.model_loader import load_sam_model - from compressors.proposal.core import generate_proposals_batch +def _(meta): + proposals = meta.get("masks", []) + return (proposals,) - sam2_processor, sam2_model = load_sam_model( - model_name="facebook/sam2.1-hiera-large" + +@app.cell +def _(Image, ImageDraw, filtered_items, image, mo, np, proposals): + from compressors.filter import ( + MaskScoringConfig, + compute_mask_features, + score_mask, + should_reject_mask, ) - input_boxes = [[box for box, _score, _text_label in filtered_items]] - proposal_batch = generate_proposals_batch( - model=sam2_model, - processor=sam2_processor, - images=[image], - bboxes_per_image=input_boxes, - ) - proposals = proposal_batch[0] if proposal_batch else [] + image_shape = (image.height, image.width) + config = MaskScoringConfig() - base_rgba = image.convert("RGBA") - _vis_image = base_rgba.copy() - summary_lines = [] + _kept = [] + _rejected = [] + for _idx, proposal in enumerate(proposals): + _feat = compute_mask_features(proposal, image_shape) + _is_rejected = should_reject_mask(_feat, config) + _score = score_mask(proposal, image_shape, config) + _owl_label = ( + filtered_items[_idx][2] if _idx < len(filtered_items) else f"obj_{_idx}" + ) + _owl_score = filtered_items[_idx][1] if _idx < len(filtered_items) else 0.0 + _owl_bbox = ( + filtered_items[_idx][0] if _idx < len(filtered_items) else [0, 0, 0, 0] + ) - colors = [ + _entry = { + "idx": _idx, + "proposal": proposal, + "features": _feat, + "mask_score": _score, + "owl_label": _owl_label, + "owl_score": _owl_score, + "owl_bbox": _owl_bbox, + } + if _is_rejected: + _rejected.append(_entry) + else: + _kept.append(_entry) + + _colors = [ (255, 0, 0, 90), (0, 255, 0, 90), (0, 0, 255, 90), @@ -198,54 +224,143 @@ def _(Image, ImageDraw, device, filtered_items, image, mo, np): (128, 0, 255, 90), ] - for _idx, ((_box, _score, _text_label), proposal) in enumerate( - zip(filtered_items, proposals) - ): - mask_np = proposal["segment"] - color = colors[_idx % len(colors)] + def _overlay_masks(base_img, entries, border_color, show_score=False): + """Overlay masks on image and draw bounding boxes with labels.""" + _rgba = base_img.convert("RGBA").copy() + for _e in entries: + _p = _e["proposal"] + _c = _colors[_e["idx"] % len(_colors)] + _mask_rgba = np.zeros((base_img.height, base_img.width, 4), dtype=np.uint8) + _mask_rgba[_p["segment"]] = _c + _rgba = Image.alpha_composite( + _rgba, Image.fromarray(_mask_rgba, mode="RGBA") + ) - mask_rgba = np.zeros((image.height, image.width, 4), dtype=np.uint8) - mask_rgba[mask_np] = color + _draw = ImageDraw.Draw(_rgba) + for _e in entries: + _x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]] + _mask_area = int(_e["proposal"]["area"]) + if show_score: + _label = f"{_e['owl_label']} | score={_e['mask_score']:.2f} | area={_mask_area}" + else: + _label = f"{_e['owl_label']} | area={_mask_area}" + _draw.rectangle([_x1, _y1, _x2, _y2], outline=border_color, width=3) - mask_img = Image.fromarray(mask_rgba, mode="RGBA") - _vis_image = Image.alpha_composite(_vis_image, mask_img) + try: + _tb = _draw.textbbox((_x1, _y1), _label) + except Exception: + _tb = (_x1, _y1, _x1 + 200, _y1 + 20) + _draw.rectangle( + [_tb[0], max(0, _tb[1] - 2), _tb[2] + 4, _tb[3] + 2], + fill=border_color, + ) + _draw.text((_x1 + 2, max(0, _y1)), _label, fill="white") + return _rgba - _draw = ImageDraw.Draw(_vis_image) - for (_box, _score, _text_label), proposal in zip(filtered_items, proposals): - _x1, _y1, _x2, _y2 = [float(v) for v in _box] - mask_area = int(proposal["area"]) - _label = f"{_text_label} | owl={_score:.3f} | mask_area={mask_area}" + _all_entries = _kept + _rejected + _all_entries.sort(key=lambda e: e["idx"]) - _draw.rectangle([_x1, _y1, _x2, _y2], outline=(255, 0, 0, 255), width=3) + _before_img = _overlay_masks(image, _all_entries, "red", show_score=False) + _after_img = _overlay_masks(image, _kept, (0, 180, 0), show_score=True) - try: - _tx1, _ty1, _tx2, _ty2 = _draw.textbbox((_x1, _y1), _label) - except Exception: - _tx1, _ty1, _tx2, _ty2 = _x1, _y1, _x1 + 220, _y1 + 20 + _draw_after = ImageDraw.Draw(_after_img) + for _e in _rejected: + _x1, _y1, _x2, _y2 = [float(v) for v in _e["owl_bbox"]] + _feat = _e["features"] + _reason_parts = [] + if _feat.area_ratio < config.min_area_ratio: + _reason_parts.append(f"area_ratio={_feat.area_ratio:.4f} config.max_area_ratio: + _reason_parts.append(f"area_ratio={_feat.area_ratio:.4f}>max") + _aspect = max(_feat.aspect_ratio, 1.0 / max(_feat.aspect_ratio, 1e-6)) + if _aspect > config.max_aspect_ratio: + _reason_parts.append(f"aspect={_aspect:.1f}>max") + if _feat.fill_ratio < config.min_fill_ratio_hard: + _reason_parts.append(f"fill={_feat.fill_ratio:.3f} config.max_components + and _feat.largest_component_ratio < config.min_largest_component_ratio + ): + _reason_parts.append(f"fragments={_feat.num_components}") + if _feat.touch_edge_count >= config.reject_edge_touch_count: + _reason_parts.append(f"edge_touch={_feat.touch_edge_count}") + if ( + _feat.touch_edge_count >= config.reject_large_edge_touch_count + and _feat.area_ratio > config.reject_large_edge_area_ratio + ): + _reason_parts.append("edge+large") - _draw.rectangle( - [_tx1, max(0, _ty1 - 2), _tx2 + 4, _ty2 + 2], - fill=(255, 0, 0, 220), - ) - _draw.text((_x1 + 2, max(0, _y1)), _label, fill="white") - summary_lines.append( - f"- {_text_label}: owl_score={_score:.3f}, mask_area={mask_area}" + _reason = ", ".join(_reason_parts) if _reason_parts else "unknown" + _label = f"X {_e['owl_label']}: {_reason}" + + _dash_len = 8 + _gap_len = 4 + for _seg_start in range(int(_x1), int(_x2), _dash_len + _gap_len): + _seg_end = min(_seg_start + _dash_len, int(_x2)) + _draw_after.line([(_seg_start, _y1), (_seg_end, _y1)], fill="red", width=2) + _draw_after.line([(_seg_start, _y2), (_seg_end, _y2)], fill="red", width=2) + for _seg_start in range(int(_y1), int(_y2), _dash_len + _gap_len): + _seg_end = min(_seg_start + _dash_len, int(_y2)) + _draw_after.line([(_x1, _seg_start), (_x1, _seg_end)], fill="red", width=2) + _draw_after.line([(_x2, _seg_start), (_x2, _seg_end)], fill="red", width=2) + + _draw_after.text((_x1 + 2, max(0, _y1 - 18)), _label, fill="red") + + _total = len(proposals) + _kept_count = len(_kept) + _rej_count = len(_rejected) + + _rej_detail_lines = [] + for _e in _rejected: + _f = _e["features"] + _rej_detail_lines.append( + f" - **{_e['owl_label']}** (idx={_e['idx']}): " + f"area_ratio={_f.area_ratio:.4f}, fill_ratio={_f.fill_ratio:.3f}, " + f"aspect_ratio={_f.aspect_ratio:.1f}, touch_edge={_f.touch_edge_count}, " + f"components={_f.num_components}, mask_score={_e['mask_score']:.3f}" ) - if not filtered_items: - summary_text = ( - "没有可用于分割的检测框,请先降低 OWLv2 的 score_threshold 或检查检测结果。" + _kept_detail_lines = [] + for _e in _kept: + _f = _e["features"] + _kept_detail_lines.append( + f" - **{_e['owl_label']}** (idx={_e['idx']}): " + f"area_ratio={_f.area_ratio:.4f}, mask_score={_e['mask_score']:.3f}" ) - elif not summary_lines: - summary_text = "没有生成任何 mask" - else: - summary_text = "\n".join(summary_lines) + + _detail_parts = [] + if _kept_detail_lines: + _detail_parts.append("**保留的 mask:**\n" + "\n".join(_kept_detail_lines)) + if _rej_detail_lines: + _detail_parts.append("**过滤掉的 mask:**\n" + "\n".join(_rej_detail_lines)) + + _detail_text = "\n\n".join(_detail_parts) if _detail_parts else "无 mask 数据" mo.vstack( [ - mo.md(f"## SAM2 分割可视化结果\n\ndevice: `{device}`"), - mo.image(_vis_image, width=700), - mo.md(summary_text), + mo.md( + "## Mask 过滤对比" + f"\n\n共 {_total} 个 mask → 保留 **{_kept_count}** 个,过滤掉 **{_rej_count}** 个" + ), + mo.hstack( + [ + mo.vstack( + [ + mo.md(f"### 过滤前({_total} 个)"), + mo.image(_before_img, width=480), + ] + ), + mo.vstack( + [ + mo.md( + f"### 过滤后({_kept_count} 个保留,{_rej_count} 个过滤)" + ), + mo.image(_after_img, width=480), + ] + ), + ] + ), + mo.md(_detail_text), ] ) return diff --git a/notebooks/verification.py b/notebooks/verification.py index 77583bc..b1e6fe6 100644 --- a/notebooks/verification.py +++ b/notebooks/verification.py @@ -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())