Spaces:
Running
Running
| """Frozen Grounding DINO and SAM2 target-category inspection.""" | |
| from __future__ import annotations | |
| from collections.abc import Callable | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import numpy as np | |
| import yaml | |
| Detection = tuple[str, float, tuple[float, float, float, float]] | |
| Detector = Callable[[np.ndarray, tuple[str, ...]], tuple[Detection, ...]] | |
| Segmenter = Callable[[np.ndarray, tuple[float, float, float, float]], np.ndarray] | |
| class InstanceDetection: | |
| category: str | |
| confidence: float | |
| mask: np.ndarray | |
| def detect_instances(detector: Detector, segmenter: Segmenter, | |
| rgb: np.ndarray, category: str) -> tuple[InstanceDetection, ...]: | |
| image = np.asarray(rgb)[..., :3] | |
| found = [] | |
| for label, score, box in detector(image, (category,)): | |
| if label != category: | |
| continue | |
| mask = np.asarray(segmenter(image, box), dtype=bool) | |
| if mask.shape == image.shape[:2] and mask.any(): | |
| found.append(InstanceDetection(label, float(score), mask)) | |
| return tuple(found) | |
| def detect_category( | |
| detector: Detector, segmenter: Segmenter, rgb: np.ndarray, category: str | |
| ) -> bool: | |
| """Use only the rendered RGB image and public target category for STOP.""" | |
| return bool(detect_instances(detector, segmenter, rgb, category)) | |
| class GroundedSAMInspector: | |
| """Frozen public models; the private semantic image never enters these heads.""" | |
| def __init__(self, config_path: Path, *, dino_model: str, sam_model: str) -> None: | |
| config = yaml.safe_load(config_path.read_text(encoding="utf-8")) | |
| self.detector, self.segmenter = _load_models( | |
| dino_model, sam_model, config | |
| ) | |
| def __call__(self, rgb: np.ndarray, _depth: np.ndarray, category: str) -> bool: | |
| return detect_category(self.detector, self.segmenter, rgb, category) | |
| def detect_instances(self, rgb: np.ndarray, category: str) -> tuple[InstanceDetection, ...]: | |
| return detect_instances(self.detector, self.segmenter, rgb, category) | |
| def close(self) -> None: | |
| del self.detector, self.segmenter | |
| def _load_models(dino_model: str, sam_model: str, config: dict): | |
| import torch | |
| from PIL import Image | |
| from transformers import ( | |
| AutoModelForZeroShotObjectDetection, | |
| AutoProcessor, | |
| Sam2Model, | |
| Sam2Processor, | |
| ) | |
| device = config["device"] | |
| dino_revision = config["grounding_dino"]["revision"] | |
| sam_revision = config["sam"]["revision"] | |
| dino_kwargs = {"revision": dino_revision} if dino_model == config["grounding_dino"]["model_id"] else {} | |
| sam_kwargs = {"revision": sam_revision} if sam_model == config["sam"]["model_id"] else {} | |
| dino_processor = AutoProcessor.from_pretrained(dino_model, **dino_kwargs) | |
| dino = AutoModelForZeroShotObjectDetection.from_pretrained( | |
| dino_model, **dino_kwargs | |
| ).to(device).eval() | |
| sam_processor = Sam2Processor.from_pretrained(sam_model, **sam_kwargs) | |
| sam = Sam2Model.from_pretrained(sam_model, **sam_kwargs).to(device).eval() | |
| def detector(rgb: np.ndarray, categories: tuple[str, ...]) -> tuple[Detection, ...]: | |
| image = Image.fromarray(rgb) | |
| inputs = dino_processor( | |
| images=image, text=[list(categories)], return_tensors="pt" | |
| ).to(device) | |
| with torch.inference_mode(): | |
| outputs = dino(**inputs) | |
| processed = dino_processor.post_process_grounded_object_detection( | |
| outputs, | |
| inputs.input_ids, | |
| threshold=config["grounding_dino"]["box_threshold"], | |
| text_threshold=config["grounding_dino"]["text_threshold"], | |
| target_sizes=[image.size[::-1]], | |
| )[0] | |
| labels = processed.get("text_labels", processed.get("labels", ())) | |
| return tuple( | |
| (str(label), float(score), tuple(float(value) for value in box)) | |
| for label, score, box in zip( | |
| labels, processed["scores"], processed["boxes"], strict=True | |
| ) | |
| ) | |
| def segmenter(rgb: np.ndarray, box: tuple[float, float, float, float]) -> np.ndarray: | |
| image = Image.fromarray(rgb) | |
| inputs = sam_processor( | |
| images=image, input_boxes=[[[*box]]], return_tensors="pt" | |
| ).to(device) | |
| with torch.inference_mode(): | |
| outputs = sam(**inputs) | |
| masks = sam_processor.post_process_masks( | |
| outputs.pred_masks.cpu(), inputs["original_sizes"] | |
| )[0] | |
| scores = outputs.iou_scores[0, 0].cpu() | |
| return masks[0, int(scores.argmax())].numpy().astype(bool) | |
| return detector, segmenter | |