SahilCodevally/codevally-vision-language-action

1"""2Visualization Module — Annotates images with detection results3and action plan overlays using OpenCV.4 5Provides two visualization modes:6    1. **Detection overlay** — All detected objects with labels + confidence.7    2. **Action overlay** — Source/destination bounding boxes, arrow, and steps.8"""9 10import logging11from typing import Any12 13import cv214import numpy as np15 16logger = logging.getLogger(__name__)17 18# ---------------------------------------------------------------------------19# Color palette (BGR for OpenCV)20# ---------------------------------------------------------------------------21COLOR_GREEN = (0, 200, 0)22COLOR_RED = (0, 0, 220)23COLOR_YELLOW = (0, 220, 220)24COLOR_CYAN = (220, 200, 0)25COLOR_WHITE = (255, 255, 255)26COLOR_BLACK = (0, 0, 0)27COLOR_ORANGE = (0, 140, 255)28COLOR_BLUE = (255, 100, 0)29 30# Distinct palette for detection overlay31_DETECTION_COLORS = [32    COLOR_GREEN,33    COLOR_CYAN,34    COLOR_YELLOW,35    COLOR_ORANGE,36    COLOR_BLUE,37    COLOR_RED,38    (200, 0, 200),  # magenta39    (100, 200, 100),  # light green40]41 42 43class ActionVisualizer:44    """45    Draws visual annotations on images for detection results and action plans.46    """47 48    # ------------------------------------------------------------------49    # Detection Overlay50    # ------------------------------------------------------------------51 52    def draw_detections(53        self,54        image: np.ndarray,55        detections: list[dict[str, Any]],56    ) -> np.ndarray:57        """58        Draw all detected objects on the image with labels and confidence.59 60        Args:61            image: Input image (H, W, C) as NumPy array.62            detections: List of detection dicts from the vision module.63 64        Returns:65            Annotated copy of the image.66        """67        if image is None or not isinstance(image, np.ndarray):68            logger.error("Invalid image for detection visualization.")69            return image if image is not None else np.zeros((400, 600, 3), dtype=np.uint8)70 71        annotated = image.copy()72 73        if not detections:74            logger.info("No detections to visualize.")75            self._put_centered_text(annotated, "No objects detected")76            return annotated77 78        for idx, det in enumerate(detections):79            color = _DETECTION_COLORS[idx % len(_DETECTION_COLORS)]80            bbox = det.get("bbox", [])81            label = det.get("label", "unknown")82            confidence = det.get("confidence", 0.0)83 84            if len(bbox) != 4:85                continue86 87            x1, y1, x2, y2 = [int(c) for c in bbox]88 89            # Draw bounding box90            cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)91 92            # Label with confidence93            text = f"{label} ({confidence:.2f})"94            text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]95 96            # Background rectangle for text97            cv2.rectangle(98                annotated,99                (x1, y1 - text_size[1] - 10),100                (x1 + text_size[0] + 6, y1),101                color,102                -1,103            )104            cv2.putText(105                annotated,106                text,107                (x1 + 3, y1 - 5),108                cv2.FONT_HERSHEY_SIMPLEX,109                0.6,110                COLOR_WHITE,111                2,112            )113 114        logger.info("Detection overlay drawn with %d objects.", len(detections))115        return annotated116 117    # ------------------------------------------------------------------118    # Action Overlay119    # ------------------------------------------------------------------120 121    def visualize(122        self,123        image: np.ndarray,124        action_plan: dict[str, Any],125    ) -> np.ndarray:126        """127        Draw the action plan overlay: source/destination boxes, arrow,128        and step labels.129 130        Args:131            image: Input image.132            action_plan: Dict from the ActionPlanner containing133                ``source_bbox``, ``dest_bbox``, ``steps``, etc.134 135        Returns:136            Annotated copy of the image.137        """138        if image is None or not isinstance(image, np.ndarray):139            logger.error("Invalid image for action visualization.")140            return image if image is not None else np.zeros((400, 600, 3), dtype=np.uint8)141 142        annotated = image.copy()143 144        source_bbox = action_plan.get("source_bbox")145        dest_bbox = action_plan.get("dest_bbox")146        steps = action_plan.get("steps", [])147        status = action_plan.get("status", "unknown")148 149        # Draw source bounding box (green)150        if source_bbox and len(source_bbox) == 4:151            self._draw_highlighted_box(152                annotated, source_bbox, COLOR_GREEN, "SOURCE"153            )154 155        # Draw destination bounding box (red)156        if dest_bbox and len(dest_bbox) == 4:157            self._draw_highlighted_box(158                annotated, dest_bbox, COLOR_RED, "DESTINATION"159            )160 161        # Draw arrow from source to destination162        if source_bbox and dest_bbox:163            self._draw_action_arrow(annotated, source_bbox, dest_bbox)164 165        # Draw step annotations166        self._draw_steps_overlay(annotated, steps, status)167 168        logger.info("Action visualization drawn (status=%s).", status)169        return annotated170 171    # ------------------------------------------------------------------172    # Drawing helpers173    # ------------------------------------------------------------------174 175    @staticmethod176    def _draw_highlighted_box(177        image: np.ndarray,178        bbox: list[int],179        color: tuple,180        label: str,181    ) -> None:182        """Draw a highlighted bounding box with a corner-style frame."""183        x1, y1, x2, y2 = [int(c) for c in bbox]184        thickness = 3185 186        # Semi-transparent overlay187        overlay = image.copy()188        cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1)189        cv2.addWeighted(overlay, 0.15, image, 0.85, 0, image)190 191        # Solid border192        cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)193 194        # Corner accents (L-shaped)195        corner_len = min(30, (x2 - x1) // 4, (y2 - y1) // 4)196        ct = thickness + 1197        # Top-left198        cv2.line(image, (x1, y1), (x1 + corner_len, y1), color, ct)199        cv2.line(image, (x1, y1), (x1, y1 + corner_len), color, ct)200        # Top-right201        cv2.line(image, (x2, y1), (x2 - corner_len, y1), color, ct)202        cv2.line(image, (x2, y1), (x2, y1 + corner_len), color, ct)203        # Bottom-left204        cv2.line(image, (x1, y2), (x1 + corner_len, y2), color, ct)205        cv2.line(image, (x1, y2), (x1, y2 - corner_len), color, ct)206        # Bottom-right207        cv2.line(image, (x2, y2), (x2 - corner_len, y2), color, ct)208        cv2.line(image, (x2, y2), (x2, y2 - corner_len), color, ct)209 210        # Label above box211        text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)[0]212        cv2.rectangle(213            image,214            (x1, y1 - text_size[1] - 14),215            (x1 + text_size[0] + 10, y1 - 2),216            color,217            -1,218        )219        cv2.putText(220            image,221            label,222            (x1 + 5, y1 - 7),223            cv2.FONT_HERSHEY_SIMPLEX,224            0.7,225            COLOR_WHITE,226            2,227        )228 229    @staticmethod230    def _draw_action_arrow(231        image: np.ndarray,232        source_bbox: list[int],233        dest_bbox: list[int],234    ) -> None:235        """Draw a curved arrow from the source centre to the destination centre."""236        sx = (source_bbox[0] + source_bbox[2]) // 2237        sy = (source_bbox[1] + source_bbox[3]) // 2238        dx = (dest_bbox[0] + dest_bbox[2]) // 2239        dy = (dest_bbox[1] + dest_bbox[3]) // 2240 241        # Draw a dashed path then a solid arrowhead242        mid_x = (sx + dx) // 2243        mid_y = min(sy, dy) - 40  # curve above244 245        # Bezier-approximated curve via polyline246        pts: list[tuple[int, int]] = []247        for t_i in range(21):248            t = t_i / 20.0249            px = int((1 - t) ** 2 * sx + 2 * (1 - t) * t * mid_x + t ** 2 * dx)250            py = int((1 - t) ** 2 * sy + 2 * (1 - t) * t * mid_y + t ** 2 * dy)251            pts.append((px, py))252 253        for i in range(len(pts) - 1):254            # Dashed effect: draw every other segment255            if i % 2 == 0:256                cv2.line(image, pts[i], pts[i + 1], COLOR_ORANGE, 3)257 258        # Arrowhead at destination259        cv2.arrowedLine(260            image, pts[-3], (dx, dy), COLOR_ORANGE, 3, tipLength=0.3261        )262 263        # "ACTION" label at midpoint264        cv2.putText(265            image,266            "ACTION",267            (mid_x - 30, mid_y - 10),268            cv2.FONT_HERSHEY_SIMPLEX,269            0.6,270            COLOR_ORANGE,271            2,272        )273 274    @staticmethod275    def _draw_steps_overlay(276        image: np.ndarray,277        steps: list[str],278        status: str,279    ) -> None:280        """Draw the steps list as a semi-transparent panel at the bottom."""281        if not steps:282            return283 284        h, w = image.shape[:2]285        line_height = 28286        panel_height = len(steps) * line_height + 20287        panel_top = h - panel_height288 289        # Semi-transparent dark panel290        overlay = image.copy()291        cv2.rectangle(overlay, (0, panel_top), (w, h), COLOR_BLACK, -1)292        cv2.addWeighted(overlay, 0.7, image, 0.3, 0, image)293 294        # Status badge295        status_color = (296            COLOR_GREEN297            if status == "success"298            else COLOR_YELLOW if status == "partial" else COLOR_RED299        )300        cv2.putText(301            image,302            f"Status: {status.upper()}",303            (10, panel_top + 20),304            cv2.FONT_HERSHEY_SIMPLEX,305            0.5,306            status_color,307            2,308        )309 310        # Steps text311        for i, step in enumerate(steps):312            y = panel_top + 20 + (i + 1) * line_height313            # Truncate long steps314            display_step = step if len(step) < 80 else step[:77] + "..."315            cv2.putText(316                image,317                display_step,318                (10, y),319                cv2.FONT_HERSHEY_SIMPLEX,320                0.5,321                COLOR_WHITE,322                1,323            )324 325    @staticmethod326    def _put_centered_text(327        image: np.ndarray,328        text: str,329        color: tuple = COLOR_WHITE,330    ) -> None:331        """Put centred text on the image."""332        h, w = image.shape[:2]333        text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 1.0, 2)[0]334        x = (w - text_size[0]) // 2335        y = (h + text_size[1]) // 2336        cv2.putText(image, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)337