codevally-vision-language-action

Running

codevally-vision-language-action / src /visualization /visualizer.py

Init commit

e1c010e 8 months ago

10.7 kB

"""
Visualization Module — Annotates images with detection results
and action plan overlays using OpenCV.
Provides two visualization modes:
    1. **Detection overlay** — All detected objects with labels + confidence.
    2. **Action overlay** — Source/destination bounding boxes, arrow, and steps.
"""

import logging
from typing import Any

import cv2
import numpy as np

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Color palette (BGR for OpenCV)
# ---------------------------------------------------------------------------
COLOR_GREEN = (0, 200, 0)
COLOR_RED = (0, 0, 220)
COLOR_YELLOW = (0, 220, 220)
COLOR_CYAN = (220, 200, 0)
COLOR_WHITE = (255, 255, 255)
COLOR_BLACK = (0, 0, 0)
COLOR_ORANGE = (0, 140, 255)
COLOR_BLUE = (255, 100, 0)

# Distinct palette for detection overlay
_DETECTION_COLORS = [
    COLOR_GREEN,
    COLOR_CYAN,
    COLOR_YELLOW,
    COLOR_ORANGE,
    COLOR_BLUE,
    COLOR_RED,
    (200, 0, 200),  # magenta
    (100, 200, 100),  # light green
]


class ActionVisualizer:
    """
    Draws visual annotations on images for detection results and action plans.
    """

    # ------------------------------------------------------------------
    # Detection Overlay
    # ------------------------------------------------------------------

    def draw_detections(
        self,
        image: np.ndarray,
        detections: list[dict[str, Any]],
    ) -> np.ndarray:
        """
        Draw all detected objects on the image with labels and confidence.
        Args:
            image: Input image (H, W, C) as NumPy array.
            detections: List of detection dicts from the vision module.
        Returns:
            Annotated copy of the image.
        """
        if image is None or not isinstance(image, np.ndarray):
            logger.error("Invalid image for detection visualization.")
            return image if image is not None else np.zeros((400, 600, 3), dtype=np.uint8)

        annotated = image.copy()

        if not detections:
            logger.info("No detections to visualize.")
            self._put_centered_text(annotated, "No objects detected")
            return annotated

        for idx, det in enumerate(detections):
            color = _DETECTION_COLORS[idx % len(_DETECTION_COLORS)]
            bbox = det.get("bbox", [])
            label = det.get("label", "unknown")
            confidence = det.get("confidence", 0.0)

            if len(bbox) != 4:
                continue

            x1, y1, x2, y2 = [int(c) for c in bbox]

            # Draw bounding box
            cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

            # Label with confidence
            text = f"{label} ({confidence:.2f})"
            text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]

            # Background rectangle for text
            cv2.rectangle(
                annotated,
                (x1, y1 - text_size[1] - 10),
                (x1 + text_size[0] + 6, y1),
                color,
                -1,
            )
            cv2.putText(
                annotated,
                text,
                (x1 + 3, y1 - 5),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                COLOR_WHITE,
                2,
            )

        logger.info("Detection overlay drawn with %d objects.", len(detections))
        return annotated

    # ------------------------------------------------------------------
    # Action Overlay
    # ------------------------------------------------------------------

    def visualize(
        self,
        image: np.ndarray,
        action_plan: dict[str, Any],
    ) -> np.ndarray:
        """
        Draw the action plan overlay: source/destination boxes, arrow,
        and step labels.
        Args:
            image: Input image.
            action_plan: Dict from the ActionPlanner containing
                ``source_bbox``, ``dest_bbox``, ``steps``, etc.
        Returns:
            Annotated copy of the image.
        """
        if image is None or not isinstance(image, np.ndarray):
            logger.error("Invalid image for action visualization.")
            return image if image is not None else np.zeros((400, 600, 3), dtype=np.uint8)

        annotated = image.copy()

        source_bbox = action_plan.get("source_bbox")
        dest_bbox = action_plan.get("dest_bbox")
        steps = action_plan.get("steps", [])
        status = action_plan.get("status", "unknown")

        # Draw source bounding box (green)
        if source_bbox and len(source_bbox) == 4:
            self._draw_highlighted_box(
                annotated, source_bbox, COLOR_GREEN, "SOURCE"
            )

        # Draw destination bounding box (red)
        if dest_bbox and len(dest_bbox) == 4:
            self._draw_highlighted_box(
                annotated, dest_bbox, COLOR_RED, "DESTINATION"
            )

        # Draw arrow from source to destination
        if source_bbox and dest_bbox:
            self._draw_action_arrow(annotated, source_bbox, dest_bbox)

        # Draw step annotations
        self._draw_steps_overlay(annotated, steps, status)

        logger.info("Action visualization drawn (status=%s).", status)
        return annotated

    # ------------------------------------------------------------------
    # Drawing helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _draw_highlighted_box(
        image: np.ndarray,
        bbox: list[int],
        color: tuple,
        label: str,
    ) -> None:
        """Draw a highlighted bounding box with a corner-style frame."""
        x1, y1, x2, y2 = [int(c) for c in bbox]
        thickness = 3

        # Semi-transparent overlay
        overlay = image.copy()
        cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1)
        cv2.addWeighted(overlay, 0.15, image, 0.85, 0, image)

        # Solid border
        cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)

        # Corner accents (L-shaped)
        corner_len = min(30, (x2 - x1) // 4, (y2 - y1) // 4)
        ct = thickness + 1
        # Top-left
        cv2.line(image, (x1, y1), (x1 + corner_len, y1), color, ct)
        cv2.line(image, (x1, y1), (x1, y1 + corner_len), color, ct)
        # Top-right
        cv2.line(image, (x2, y1), (x2 - corner_len, y1), color, ct)
        cv2.line(image, (x2, y1), (x2, y1 + corner_len), color, ct)
        # Bottom-left
        cv2.line(image, (x1, y2), (x1 + corner_len, y2), color, ct)
        cv2.line(image, (x1, y2), (x1, y2 - corner_len), color, ct)
        # Bottom-right
        cv2.line(image, (x2, y2), (x2 - corner_len, y2), color, ct)
        cv2.line(image, (x2, y2), (x2, y2 - corner_len), color, ct)

        # Label above box
        text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)[0]
        cv2.rectangle(
            image,
            (x1, y1 - text_size[1] - 14),
            (x1 + text_size[0] + 10, y1 - 2),
            color,
            -1,
        )
        cv2.putText(
            image,
            label,
            (x1 + 5, y1 - 7),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            COLOR_WHITE,
            2,
        )

    @staticmethod
    def _draw_action_arrow(
        image: np.ndarray,
        source_bbox: list[int],
        dest_bbox: list[int],
    ) -> None:
        """Draw a curved arrow from the source centre to the destination centre."""
        sx = (source_bbox[0] + source_bbox[2]) // 2
        sy = (source_bbox[1] + source_bbox[3]) // 2
        dx = (dest_bbox[0] + dest_bbox[2]) // 2
        dy = (dest_bbox[1] + dest_bbox[3]) // 2

        # Draw a dashed path then a solid arrowhead
        mid_x = (sx + dx) // 2
        mid_y = min(sy, dy) - 40  # curve above

        # Bezier-approximated curve via polyline
        pts: list[tuple[int, int]] = []
        for t_i in range(21):
            t = t_i / 20.0
            px = int((1 - t) ** 2 * sx + 2 * (1 - t) * t * mid_x + t ** 2 * dx)
            py = int((1 - t) ** 2 * sy + 2 * (1 - t) * t * mid_y + t ** 2 * dy)
            pts.append((px, py))

        for i in range(len(pts) - 1):
            # Dashed effect: draw every other segment
            if i % 2 == 0:
                cv2.line(image, pts[i], pts[i + 1], COLOR_ORANGE, 3)

        # Arrowhead at destination
        cv2.arrowedLine(
            image, pts[-3], (dx, dy), COLOR_ORANGE, 3, tipLength=0.3
        )

        # "ACTION" label at midpoint
        cv2.putText(
            image,
            "ACTION",
            (mid_x - 30, mid_y - 10),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.6,
            COLOR_ORANGE,
            2,
        )

    @staticmethod
    def _draw_steps_overlay(
        image: np.ndarray,
        steps: list[str],
        status: str,
    ) -> None:
        """Draw the steps list as a semi-transparent panel at the bottom."""
        if not steps:
            return

        h, w = image.shape[:2]
        line_height = 28
        panel_height = len(steps) * line_height + 20
        panel_top = h - panel_height

        # Semi-transparent dark panel
        overlay = image.copy()
        cv2.rectangle(overlay, (0, panel_top), (w, h), COLOR_BLACK, -1)
        cv2.addWeighted(overlay, 0.7, image, 0.3, 0, image)

        # Status badge
        status_color = (
            COLOR_GREEN
            if status == "success"
            else COLOR_YELLOW if status == "partial" else COLOR_RED
        )
        cv2.putText(
            image,
            f"Status: {status.upper()}",
            (10, panel_top + 20),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            status_color,
            2,
        )

        # Steps text
        for i, step in enumerate(steps):
            y = panel_top + 20 + (i + 1) * line_height
            # Truncate long steps
            display_step = step if len(step) < 80 else step[:77] + "..."
            cv2.putText(
                image,
                display_step,
                (10, y),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                COLOR_WHITE,
                1,
            )

    @staticmethod
    def _put_centered_text(
        image: np.ndarray,
        text: str,
        color: tuple = COLOR_WHITE,
    ) -> None:
        """Put centred text on the image."""
        h, w = image.shape[:2]
        text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 1.0, 2)[0]
        x = (w - text_size[0]) // 2
        y = (h + text_size[1]) // 2
        cv2.putText(image, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)

	"""
	Visualization Module — Annotates images with detection results
	and action plan overlays using OpenCV.

	Provides two visualization modes:
	1. Detection overlay — All detected objects with labels + confidence.
	2. Action overlay — Source/destination bounding boxes, arrow, and steps.
	"""

	import logging
	from typing import Any

	import cv2
	import numpy as np

	logger = logging.getLogger(__name__)

	# ---------------------------------------------------------------------------
	# Color palette (BGR for OpenCV)
	# ---------------------------------------------------------------------------
	COLOR_GREEN = (0, 200, 0)
	COLOR_RED = (0, 0, 220)
	COLOR_YELLOW = (0, 220, 220)
	COLOR_CYAN = (220, 200, 0)
	COLOR_WHITE = (255, 255, 255)
	COLOR_BLACK = (0, 0, 0)
	COLOR_ORANGE = (0, 140, 255)
	COLOR_BLUE = (255, 100, 0)

	# Distinct palette for detection overlay
	_DETECTION_COLORS = [
	COLOR_GREEN,
	COLOR_CYAN,
	COLOR_YELLOW,
	COLOR_ORANGE,
	COLOR_BLUE,
	COLOR_RED,
	(200, 0, 200), # magenta
	(100, 200, 100), # light green
	]


	class ActionVisualizer:
	"""
	Draws visual annotations on images for detection results and action plans.
	"""

	# ------------------------------------------------------------------
	# Detection Overlay
	# ------------------------------------------------------------------

	def draw_detections(
	self,
	image: np.ndarray,
	detections: list[dict[str, Any]],
	) -> np.ndarray:
	"""
	Draw all detected objects on the image with labels and confidence.

	Args:
	image: Input image (H, W, C) as NumPy array.
	detections: List of detection dicts from the vision module.

	Returns:
	Annotated copy of the image.
	"""
	if image is None or not isinstance(image, np.ndarray):
	logger.error("Invalid image for detection visualization.")
	return image if image is not None else np.zeros((400, 600, 3), dtype=np.uint8)

	annotated = image.copy()

	if not detections:
	logger.info("No detections to visualize.")
	self._put_centered_text(annotated, "No objects detected")
	return annotated

	for idx, det in enumerate(detections):
	color = _DETECTION_COLORS[idx % len(_DETECTION_COLORS)]
	bbox = det.get("bbox", [])
	label = det.get("label", "unknown")
	confidence = det.get("confidence", 0.0)

	if len(bbox) != 4:
	continue

	x1, y1, x2, y2 = [int(c) for c in bbox]

	# Draw bounding box
	cv2.rectangle(annotated, (x1, y1), (x2, y2), color, 2)

	# Label with confidence
	text = f"{label} ({confidence:.2f})"
	text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]

	# Background rectangle for text
	cv2.rectangle(
	annotated,
	(x1, y1 - text_size[1] - 10),
	(x1 + text_size[0] + 6, y1),
	color,
	-1,
	)
	cv2.putText(
	annotated,
	text,
	(x1 + 3, y1 - 5),
	cv2.FONT_HERSHEY_SIMPLEX,
	0.6,
	COLOR_WHITE,
	2,
	)

	logger.info("Detection overlay drawn with %d objects.", len(detections))
	return annotated

	# ------------------------------------------------------------------
	# Action Overlay
	# ------------------------------------------------------------------

	def visualize(
	self,
	image: np.ndarray,
	action_plan: dict[str, Any],
	) -> np.ndarray:
	"""
	Draw the action plan overlay: source/destination boxes, arrow,
	and step labels.

	Args:
	image: Input image.
	action_plan: Dict from the ActionPlanner containing
	``source_bbox``, ``dest_bbox``, ``steps``, etc.

	Returns:
	Annotated copy of the image.
	"""
	if image is None or not isinstance(image, np.ndarray):
	logger.error("Invalid image for action visualization.")
	return image if image is not None else np.zeros((400, 600, 3), dtype=np.uint8)

	annotated = image.copy()

	source_bbox = action_plan.get("source_bbox")
	dest_bbox = action_plan.get("dest_bbox")
	steps = action_plan.get("steps", [])
	status = action_plan.get("status", "unknown")

	# Draw source bounding box (green)
	if source_bbox and len(source_bbox) == 4:
	self._draw_highlighted_box(
	annotated, source_bbox, COLOR_GREEN, "SOURCE"
	)

	# Draw destination bounding box (red)
	if dest_bbox and len(dest_bbox) == 4:
	self._draw_highlighted_box(
	annotated, dest_bbox, COLOR_RED, "DESTINATION"
	)

	# Draw arrow from source to destination
	if source_bbox and dest_bbox:
	self._draw_action_arrow(annotated, source_bbox, dest_bbox)

	# Draw step annotations
	self._draw_steps_overlay(annotated, steps, status)

	logger.info("Action visualization drawn (status=%s).", status)
	return annotated

	# ------------------------------------------------------------------
	# Drawing helpers
	# ------------------------------------------------------------------

	@staticmethod
	def _draw_highlighted_box(
	image: np.ndarray,
	bbox: list[int],
	color: tuple,
	label: str,
	) -> None:
	"""Draw a highlighted bounding box with a corner-style frame."""
	x1, y1, x2, y2 = [int(c) for c in bbox]
	thickness = 3

	# Semi-transparent overlay
	overlay = image.copy()
	cv2.rectangle(overlay, (x1, y1), (x2, y2), color, -1)
	cv2.addWeighted(overlay, 0.15, image, 0.85, 0, image)

	# Solid border
	cv2.rectangle(image, (x1, y1), (x2, y2), color, thickness)

	# Corner accents (L-shaped)
	corner_len = min(30, (x2 - x1) // 4, (y2 - y1) // 4)
	ct = thickness + 1
	# Top-left
	cv2.line(image, (x1, y1), (x1 + corner_len, y1), color, ct)
	cv2.line(image, (x1, y1), (x1, y1 + corner_len), color, ct)
	# Top-right
	cv2.line(image, (x2, y1), (x2 - corner_len, y1), color, ct)
	cv2.line(image, (x2, y1), (x2, y1 + corner_len), color, ct)
	# Bottom-left
	cv2.line(image, (x1, y2), (x1 + corner_len, y2), color, ct)
	cv2.line(image, (x1, y2), (x1, y2 - corner_len), color, ct)
	# Bottom-right
	cv2.line(image, (x2, y2), (x2 - corner_len, y2), color, ct)
	cv2.line(image, (x2, y2), (x2, y2 - corner_len), color, ct)

	# Label above box
	text_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)[0]
	cv2.rectangle(
	image,
	(x1, y1 - text_size[1] - 14),
	(x1 + text_size[0] + 10, y1 - 2),
	color,
	-1,
	)
	cv2.putText(
	image,
	label,
	(x1 + 5, y1 - 7),
	cv2.FONT_HERSHEY_SIMPLEX,
	0.7,
	COLOR_WHITE,
	2,
	)

	@staticmethod
	def _draw_action_arrow(
	image: np.ndarray,
	source_bbox: list[int],
	dest_bbox: list[int],
	) -> None:
	"""Draw a curved arrow from the source centre to the destination centre."""
	sx = (source_bbox[0] + source_bbox[2]) // 2
	sy = (source_bbox[1] + source_bbox[3]) // 2
	dx = (dest_bbox[0] + dest_bbox[2]) // 2
	dy = (dest_bbox[1] + dest_bbox[3]) // 2

	# Draw a dashed path then a solid arrowhead
	mid_x = (sx + dx) // 2
	mid_y = min(sy, dy) - 40 # curve above

	# Bezier-approximated curve via polyline
	pts: list[tuple[int, int]] = []
	for t_i in range(21):
	t = t_i / 20.0
	px = int((1 - t) ** 2 * sx + 2 * (1 - t) * t * mid_x + t ** 2 * dx)
	py = int((1 - t) ** 2 * sy + 2 * (1 - t) * t * mid_y + t ** 2 * dy)
	pts.append((px, py))

	for i in range(len(pts) - 1):
	# Dashed effect: draw every other segment
	if i % 2 == 0:
	cv2.line(image, pts[i], pts[i + 1], COLOR_ORANGE, 3)

	# Arrowhead at destination
	cv2.arrowedLine(
	image, pts[-3], (dx, dy), COLOR_ORANGE, 3, tipLength=0.3
	)

	# "ACTION" label at midpoint
	cv2.putText(
	image,
	"ACTION",
	(mid_x - 30, mid_y - 10),
	cv2.FONT_HERSHEY_SIMPLEX,
	0.6,
	COLOR_ORANGE,
	2,
	)

	@staticmethod
	def _draw_steps_overlay(
	image: np.ndarray,
	steps: list[str],
	status: str,
	) -> None:
	"""Draw the steps list as a semi-transparent panel at the bottom."""
	if not steps:
	return

	h, w = image.shape[:2]
	line_height = 28
	panel_height = len(steps) * line_height + 20
	panel_top = h - panel_height

	# Semi-transparent dark panel
	overlay = image.copy()
	cv2.rectangle(overlay, (0, panel_top), (w, h), COLOR_BLACK, -1)
	cv2.addWeighted(overlay, 0.7, image, 0.3, 0, image)

	# Status badge
	status_color = (
	COLOR_GREEN
	if status == "success"
	else COLOR_YELLOW if status == "partial" else COLOR_RED
	)
	cv2.putText(
	image,
	f"Status: {status.upper()}",
	(10, panel_top + 20),
	cv2.FONT_HERSHEY_SIMPLEX,
	0.5,
	status_color,
	2,
	)

	# Steps text
	for i, step in enumerate(steps):
	y = panel_top + 20 + (i + 1) * line_height
	# Truncate long steps
	display_step = step if len(step) < 80 else step[:77] + "..."
	cv2.putText(
	image,
	display_step,
	(10, y),
	cv2.FONT_HERSHEY_SIMPLEX,
	0.5,
	COLOR_WHITE,
	1,
	)

	@staticmethod
	def _put_centered_text(
	image: np.ndarray,
	text: str,
	color: tuple = COLOR_WHITE,
	) -> None:
	"""Put centred text on the image."""
	h, w = image.shape[:2]
	text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 1.0, 2)[0]
	x = (w - text_size[0]) // 2
	y = (h + text_size[1]) // 2
	cv2.putText(image, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2)