codevally-vision-language-action

Running

codevally-vision-language-action / src /vision /grounding_dino_detector.py

Init commit

e1c010e 8 months ago

12.7 kB

"""
Grounding DINO Detector — Open-vocabulary object detection using
text prompts via HuggingFace Transformers.
Unlike YOLO (fixed 80 COCO classes), Grounding DINO can detect *any*
object described in natural language. This makes it ideal for demos
where users reference domain-specific objects like "defective board",
"red bin", or "toolbox".
"""

import logging
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)


class GroundingDINODetector:
    """
    Detects objects in images using text-prompted Grounding DINO.
    The model is loaded once in the constructor and reused. Detection
    requires a text prompt describing what objects to look for.
    Attributes:
        model: The loaded Grounding DINO model.
        processor: The image/text processor.
    """

    def __init__(
        self,
        model_name: str = "IDEA-Research/grounding-dino-tiny",
        confidence: float = 0.25,
    ):
        """
        Initialize the Grounding DINO detector.
        Args:
            model_name: HuggingFace model identifier.
            confidence: Minimum confidence (box_threshold) for detections.
        """
        self.model_name = model_name
        self.confidence = confidence
        self.model = None
        self.processor = None

        self._load_model()

    def _load_model(self) -> None:
        """Load the Grounding DINO model and processor."""
        try:
            import torch
            from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection

            logger.info("Loading Grounding DINO model: %s", self.model_name)

            self.processor = AutoProcessor.from_pretrained(self.model_name)
            self.model = AutoModelForZeroShotObjectDetection.from_pretrained(
                self.model_name,
            )

            # Use GPU if available, otherwise CPU
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            self.model = self.model.to(self.device)
            self.model.eval()

            logger.info(
                "Grounding DINO loaded successfully (device=%s).", self.device
            )
        except Exception as exc:
            logger.error(
                "Failed to load Grounding DINO '%s': %s", self.model_name, exc
            )
            raise RuntimeError(
                f"Could not load Grounding DINO model '{self.model_name}'. "
                f"Ensure transformers and torch are installed."
            ) from exc

    def detect(
        self,
        image: np.ndarray,
        text_prompt: str | None = None,
    ) -> list[dict[str, Any]]:
        """
        Run open-vocabulary detection on an image.
        Args:
            image: Input image as a NumPy array (H, W, C) in RGB format.
            text_prompt: Dot-separated list of object descriptions to detect.
                         Example: ``"red bin. defective board. toolbox."``
                         If None, uses a broad default prompt.
        Returns:
            List of detected objects with label, confidence, and bbox.
        Raises:
            ValueError: If the input image is invalid.
            RuntimeError: If the model is not loaded.
        """
        if self.model is None or self.processor is None:
            raise RuntimeError("Grounding DINO model is not loaded.")

        if image is None or not isinstance(image, np.ndarray):
            raise ValueError("Invalid image: expected a NumPy ndarray.")

        if image.ndim < 2:
            raise ValueError(
                f"Invalid image dimensions: expected 2D or 3D array, "
                f"got {image.ndim}D."
            )

        if not text_prompt:
            text_prompt = "object."

        # Ensure prompt ends with a period (required by Grounding DINO)
        if not text_prompt.strip().endswith("."):
            text_prompt = text_prompt.strip() + "."

        logger.info(
            "Running Grounding DINO detection: prompt='%s', confidence=%.2f",
            text_prompt,
            self.confidence,
        )

        try:
            import torch
            from PIL import Image as PILImage

            # Convert numpy to PIL Image
            pil_image = PILImage.fromarray(image)

            # Process inputs
            inputs = self.processor(
                images=pil_image,
                text=text_prompt,
                return_tensors="pt",
            ).to(self.device)

            # Run inference
            with torch.no_grad():
                outputs = self.model(**inputs)

            # Post-process results (transformers >=5.0 API)
            results = self.processor.post_process_grounded_object_detection(
                outputs,
                threshold=self.confidence,
                text_threshold=self.confidence,
                target_sizes=[pil_image.size[::-1]],  # (height, width)
            )

            detections = self._parse_results(results)
            detections = self._apply_nms(detections)
            logger.info("Grounding DINO detected %d objects (after NMS).", len(detections))
            return detections

        except Exception as exc:
            logger.error("Grounding DINO inference failed: %s", exc)
            return []

    def _parse_results(
        self, results: list[dict[str, Any]]
    ) -> list[dict[str, Any]]:
        """Parse Grounding DINO results into structured detections."""
        detections: list[dict[str, Any]] = []

        if not results:
            return detections

        result = results[0]  # Single image
        boxes = result.get("boxes", [])
        scores = result.get("scores", [])
        # transformers >=5.0 returns text_labels (list[str]) directly
        labels = result.get("text_labels") or result.get("labels", [])

        for i in range(len(boxes)):
            try:
                bbox = boxes[i].cpu().numpy().tolist()
                confidence = float(scores[i].cpu().numpy())
                raw_label = labels[i] if i < len(labels) else "object"
                label = str(raw_label).strip()

                detections.append(
                    {
                        "label": label,
                        "confidence": float(round(confidence, 4)),
                        "bbox": [int(coord) for coord in bbox],
                    }
                )
            except (IndexError, AttributeError) as exc:
                logger.warning("Failed to parse detection %d: %s", i, exc)
                continue

        return detections

    @staticmethod
    def _iou(box_a: list[int], box_b: list[int]) -> float:
        """
        Compute Intersection-over-Union (IoU) between two bounding boxes.
        Boxes are in [x1, y1, x2, y2] format.
        """
        xa1, ya1, xa2, ya2 = box_a
        xb1, yb1, xb2, yb2 = box_b

        inter_x1 = max(xa1, xb1)
        inter_y1 = max(ya1, yb1)
        inter_x2 = min(xa2, xb2)
        inter_y2 = min(ya2, yb2)

        inter_w = max(0, inter_x2 - inter_x1)
        inter_h = max(0, inter_y2 - inter_y1)
        inter_area = inter_w * inter_h

        area_a = max(0, xa2 - xa1) * max(0, ya2 - ya1)
        area_b = max(0, xb2 - xb1) * max(0, yb2 - yb1)
        union_area = area_a + area_b - inter_area

        return inter_area / union_area if union_area > 0 else 0.0

    def _apply_nms(
        self,
        detections: list[dict[str, Any]],
        iou_threshold: float = 0.5,
    ) -> list[dict[str, Any]]:
        """
        Apply Non-Max Suppression to remove redundant detections.
        When two boxes overlap significantly (IoU >= iou_threshold), keep the
        one with the longer/more-specific label (e.g. "red bin" beats "red")
        or the higher confidence if labels have the same length.
        Args:
            detections: Raw list of detection dicts.
            iou_threshold: IoU overlap threshold above which one box is suppressed.
        Returns:
            Filtered list with duplicates removed.
        """
        if len(detections) <= 1:
            return detections

        # Sort by confidence descending
        detections = sorted(detections, key=lambda d: d["confidence"], reverse=True)
        kept: list[dict[str, Any]] = []

        for det in detections:
            suppress = False
            for kept_det in kept:
                iou = self._iou(det["bbox"], kept_det["bbox"])
                if iou >= iou_threshold:
                    # Prefer more specific (longer) label; mark current as suppressed
                    if len(det["label"]) > len(kept_det["label"]):
                        # Current is more specific — replace the kept one
                        kept.remove(kept_det)
                        break  # will be added below
                    else:
                        suppress = True
                        break

            if not suppress:
                kept.append(det)

        if len(kept) < len(detections):
            logger.info(
                "NMS removed %d redundant detections.", len(detections) - len(kept)
            )

        return kept

    @staticmethod
    def build_prompt_from_command(command: str) -> str:
        """
        Extract object-like noun phrases from a user command and build
        a Grounding DINO text prompt.
        Condition adjectives (defective, damaged, broken, etc.) are stripped
        because they describe *state*, not visual appearance, and confuse
        Grounding DINO. The underlying object noun is kept.
        Args:
            command: User's natural-language command.
        Returns:
            A dot-separated prompt string for Grounding DINO.
        Examples:
            >>> build_prompt_from_command("Move the red bin to the table")
            "red bin. table."
            >>> build_prompt_from_command("Pick up the defective board")
            "board."
        """
        import re
        cmd_lower = command.lower().strip()

        # ------------------------------------------------------------------
        # Step 1: Strip condition/state adjectives — they describe condition
        # not visual appearance so Grounding DINO can't detect them.
        # ------------------------------------------------------------------
        CONDITION_ADJECTIVES = {
            "defective", "damaged", "broken", "faulty", "worn", "old",
            "cracked", "scratched", "dented", "bent", "burnt", "corroded",
            "rusted", "leaking", "contaminated", "expired", "rejected",
            "malfunctioning", "misaligned", "missing", "incorrect",
            "closest", "nearest", "specific", "certain", "particular",
        }
        words = cmd_lower.split()
        cleaned_words = [w for w in words if w not in CONDITION_ADJECTIVES]
        cmd_lower = " ".join(cleaned_words)

        # ------------------------------------------------------------------
        # Step 2: Remove action verbs, articles, and prepositions
        # ------------------------------------------------------------------
        stop_words = {
            "the", "a", "an", "and", "or", "to", "from", "in", "on", "at",
            "into", "onto", "of", "with", "for", "it", "up", "back", "next",
            "pick", "place", "move", "inspect", "rotate", "remove", "find",
            "put", "take", "grab", "get", "bring", "carry", "push", "pull",
            "near", "close", "under", "over", "above", "below", "beside",
            "behind", "please", "can", "you", "will", "should", "must",
            "then", "also", "there", "here", "this", "that",
        }

        # ------------------------------------------------------------------
        # Step 3: Split on locational prepositions to isolate noun chunks
        # ------------------------------------------------------------------
        chunks = re.split(
            r'\b(?:to|from|into|onto|next to|near|beside|under|over|on top of'
            r'|and\s+(?:place|put|move|bring|carry|push|pull|drop|set|lay))\b',
            cmd_lower,
        )

        object_terms: list[str] = []
        seen: set[str] = set()
        for chunk in chunks:
            chunk_words = chunk.strip().split()
            cleaned = [w for w in chunk_words if w not in stop_words]
            if cleaned:
                term = " ".join(cleaned).strip()
                # Deduplicate and skip trivially short terms
                if term and len(term) > 1 and term not in seen:
                    object_terms.append(term)
                    seen.add(term)

        if not object_terms:
            return cmd_lower + "."

        prompt = ". ".join(object_terms) + "."
        logger.info("Built Grounding DINO prompt: '%s'", prompt)
        return prompt

	"""
	Grounding DINO Detector — Open-vocabulary object detection using
	text prompts via HuggingFace Transformers.

	Unlike YOLO (fixed 80 COCO classes), Grounding DINO can detect any
	object described in natural language. This makes it ideal for demos
	where users reference domain-specific objects like "defective board",
	"red bin", or "toolbox".
	"""

	import logging
	from typing import Any

	import numpy as np

	logger = logging.getLogger(__name__)


	class GroundingDINODetector:
	"""
	Detects objects in images using text-prompted Grounding DINO.

	The model is loaded once in the constructor and reused. Detection
	requires a text prompt describing what objects to look for.

	Attributes:
	model: The loaded Grounding DINO model.
	processor: The image/text processor.
	"""

	def __init__(
	self,
	model_name: str = "IDEA-Research/grounding-dino-tiny",
	confidence: float = 0.25,
	):
	"""
	Initialize the Grounding DINO detector.

	Args:
	model_name: HuggingFace model identifier.
	confidence: Minimum confidence (box_threshold) for detections.
	"""
	self.model_name = model_name
	self.confidence = confidence
	self.model = None
	self.processor = None

	self._load_model()

	def _load_model(self) -> None:
	"""Load the Grounding DINO model and processor."""
	try:
	import torch
	from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection

	logger.info("Loading Grounding DINO model: %s", self.model_name)

	self.processor = AutoProcessor.from_pretrained(self.model_name)
	self.model = AutoModelForZeroShotObjectDetection.from_pretrained(
	self.model_name,
	)

	# Use GPU if available, otherwise CPU
	self.device = "cuda" if torch.cuda.is_available() else "cpu"
	self.model = self.model.to(self.device)
	self.model.eval()

	logger.info(
	"Grounding DINO loaded successfully (device=%s).", self.device
	)
	except Exception as exc:
	logger.error(
	"Failed to load Grounding DINO '%s': %s", self.model_name, exc
	)
	raise RuntimeError(
	f"Could not load Grounding DINO model '{self.model_name}'. "
	f"Ensure transformers and torch are installed."
	) from exc

	def detect(
	self,
	image: np.ndarray,
	text_prompt: str \| None = None,
	) -> list[dict[str, Any]]:
	"""
	Run open-vocabulary detection on an image.

	Args:
	image: Input image as a NumPy array (H, W, C) in RGB format.
	text_prompt: Dot-separated list of object descriptions to detect.
	Example: ``"red bin. defective board. toolbox."``
	If None, uses a broad default prompt.

	Returns:
	List of detected objects with label, confidence, and bbox.

	Raises:
	ValueError: If the input image is invalid.
	RuntimeError: If the model is not loaded.
	"""
	if self.model is None or self.processor is None:
	raise RuntimeError("Grounding DINO model is not loaded.")

	if image is None or not isinstance(image, np.ndarray):
	raise ValueError("Invalid image: expected a NumPy ndarray.")

	if image.ndim < 2:
	raise ValueError(
	f"Invalid image dimensions: expected 2D or 3D array, "
	f"got {image.ndim}D."
	)

	if not text_prompt:
	text_prompt = "object."

	# Ensure prompt ends with a period (required by Grounding DINO)
	if not text_prompt.strip().endswith("."):
	text_prompt = text_prompt.strip() + "."

	logger.info(
	"Running Grounding DINO detection: prompt='%s', confidence=%.2f",
	text_prompt,
	self.confidence,
	)

	try:
	import torch
	from PIL import Image as PILImage

	# Convert numpy to PIL Image
	pil_image = PILImage.fromarray(image)

	# Process inputs
	inputs = self.processor(
	images=pil_image,
	text=text_prompt,
	return_tensors="pt",
	).to(self.device)

	# Run inference
	with torch.no_grad():
	outputs = self.model(**inputs)

	# Post-process results (transformers >=5.0 API)
	results = self.processor.post_process_grounded_object_detection(
	outputs,
	threshold=self.confidence,
	text_threshold=self.confidence,
	target_sizes=[pil_image.size[::-1]], # (height, width)
	)

	detections = self._parse_results(results)
	detections = self._apply_nms(detections)
	logger.info("Grounding DINO detected %d objects (after NMS).", len(detections))
	return detections

	except Exception as exc:
	logger.error("Grounding DINO inference failed: %s", exc)
	return []

	def _parse_results(
	self, results: list[dict[str, Any]]
	) -> list[dict[str, Any]]:
	"""Parse Grounding DINO results into structured detections."""
	detections: list[dict[str, Any]] = []

	if not results:
	return detections

	result = results[0] # Single image
	boxes = result.get("boxes", [])
	scores = result.get("scores", [])
	# transformers >=5.0 returns text_labels (list[str]) directly
	labels = result.get("text_labels") or result.get("labels", [])

	for i in range(len(boxes)):
	try:
	bbox = boxes[i].cpu().numpy().tolist()
	confidence = float(scores[i].cpu().numpy())
	raw_label = labels[i] if i < len(labels) else "object"
	label = str(raw_label).strip()

	detections.append(
	{
	"label": label,
	"confidence": float(round(confidence, 4)),
	"bbox": [int(coord) for coord in bbox],
	}
	)
	except (IndexError, AttributeError) as exc:
	logger.warning("Failed to parse detection %d: %s", i, exc)
	continue

	return detections

	@staticmethod
	def _iou(box_a: list[int], box_b: list[int]) -> float:
	"""
	Compute Intersection-over-Union (IoU) between two bounding boxes.

	Boxes are in [x1, y1, x2, y2] format.
	"""
	xa1, ya1, xa2, ya2 = box_a
	xb1, yb1, xb2, yb2 = box_b

	inter_x1 = max(xa1, xb1)
	inter_y1 = max(ya1, yb1)
	inter_x2 = min(xa2, xb2)
	inter_y2 = min(ya2, yb2)

	inter_w = max(0, inter_x2 - inter_x1)
	inter_h = max(0, inter_y2 - inter_y1)
	inter_area = inter_w * inter_h

	area_a = max(0, xa2 - xa1) * max(0, ya2 - ya1)
	area_b = max(0, xb2 - xb1) * max(0, yb2 - yb1)
	union_area = area_a + area_b - inter_area

	return inter_area / union_area if union_area > 0 else 0.0

	def _apply_nms(
	self,
	detections: list[dict[str, Any]],
	iou_threshold: float = 0.5,
	) -> list[dict[str, Any]]:
	"""
	Apply Non-Max Suppression to remove redundant detections.

	When two boxes overlap significantly (IoU >= iou_threshold), keep the
	one with the longer/more-specific label (e.g. "red bin" beats "red")
	or the higher confidence if labels have the same length.

	Args:
	detections: Raw list of detection dicts.
	iou_threshold: IoU overlap threshold above which one box is suppressed.

	Returns:
	Filtered list with duplicates removed.
	"""
	if len(detections) <= 1:
	return detections

	# Sort by confidence descending
	detections = sorted(detections, key=lambda d: d["confidence"], reverse=True)
	kept: list[dict[str, Any]] = []

	for det in detections:
	suppress = False
	for kept_det in kept:
	iou = self._iou(det["bbox"], kept_det["bbox"])
	if iou >= iou_threshold:
	# Prefer more specific (longer) label; mark current as suppressed
	if len(det["label"]) > len(kept_det["label"]):
	# Current is more specific — replace the kept one
	kept.remove(kept_det)
	break # will be added below
	else:
	suppress = True
	break

	if not suppress:
	kept.append(det)

	if len(kept) < len(detections):
	logger.info(
	"NMS removed %d redundant detections.", len(detections) - len(kept)
	)

	return kept

	@staticmethod
	def build_prompt_from_command(command: str) -> str:
	"""
	Extract object-like noun phrases from a user command and build
	a Grounding DINO text prompt.

	Condition adjectives (defective, damaged, broken, etc.) are stripped
	because they describe state, not visual appearance, and confuse
	Grounding DINO. The underlying object noun is kept.

	Args:
	command: User's natural-language command.

	Returns:
	A dot-separated prompt string for Grounding DINO.

	Examples:
	>>> build_prompt_from_command("Move the red bin to the table")
	"red bin. table."
	>>> build_prompt_from_command("Pick up the defective board")
	"board."
	"""
	import re
	cmd_lower = command.lower().strip()

	# ------------------------------------------------------------------
	# Step 1: Strip condition/state adjectives — they describe condition
	# not visual appearance so Grounding DINO can't detect them.
	# ------------------------------------------------------------------
	CONDITION_ADJECTIVES = {
	"defective", "damaged", "broken", "faulty", "worn", "old",
	"cracked", "scratched", "dented", "bent", "burnt", "corroded",
	"rusted", "leaking", "contaminated", "expired", "rejected",
	"malfunctioning", "misaligned", "missing", "incorrect",
	"closest", "nearest", "specific", "certain", "particular",
	}
	words = cmd_lower.split()
	cleaned_words = [w for w in words if w not in CONDITION_ADJECTIVES]
	cmd_lower = " ".join(cleaned_words)

	# ------------------------------------------------------------------
	# Step 2: Remove action verbs, articles, and prepositions
	# ------------------------------------------------------------------
	stop_words = {
	"the", "a", "an", "and", "or", "to", "from", "in", "on", "at",
	"into", "onto", "of", "with", "for", "it", "up", "back", "next",
	"pick", "place", "move", "inspect", "rotate", "remove", "find",
	"put", "take", "grab", "get", "bring", "carry", "push", "pull",
	"near", "close", "under", "over", "above", "below", "beside",
	"behind", "please", "can", "you", "will", "should", "must",
	"then", "also", "there", "here", "this", "that",
	}

	# ------------------------------------------------------------------
	# Step 3: Split on locational prepositions to isolate noun chunks
	# ------------------------------------------------------------------
	chunks = re.split(
	r'\b(?:to\|from\|into\|onto\|next to\|near\|beside\|under\|over\|on top of'
	r'\|and\s+(?:place\|put\|move\|bring\|carry\|push\|pull\|drop\|set\|lay))\b',
	cmd_lower,
	)

	object_terms: list[str] = []
	seen: set[str] = set()
	for chunk in chunks:
	chunk_words = chunk.strip().split()
	cleaned = [w for w in chunk_words if w not in stop_words]
	if cleaned:
	term = " ".join(cleaned).strip()
	# Deduplicate and skip trivially short terms
	if term and len(term) > 1 and term not in seen:
	object_terms.append(term)
	seen.add(term)

	if not object_terms:
	return cmd_lower + "."

	prompt = ". ".join(object_terms) + "."
	logger.info("Built Grounding DINO prompt: '%s'", prompt)
	return prompt