Newer
Older
RobotCar / src / pc / vision / overlay.py
"""
overlay
画像処理の結果をカメラ映像に重ねて描画するモジュール
チェックボックスで個別に ON/OFF できる
"""

from dataclasses import dataclass

import cv2
import numpy as np

from pc.vision.line_detector import LineDetectResult

# 描画色の定義 (BGR)
COLOR_LINE: tuple = (0, 255, 0)
COLOR_CENTER: tuple = (0, 255, 255)
COLOR_REGION: tuple = (255, 0, 0)

# 二値化オーバーレイの不透明度
BINARY_OPACITY: float = 0.4


@dataclass
class OverlayFlags:
    """オーバーレイ表示項目のフラグ

    Attributes:
        binary: 二値化画像の半透明表示
        detect_region: 検出領域の枠
        poly_curve: フィッティング曲線
        center_line: 画像中心線
    """
    binary: bool = False
    detect_region: bool = False
    poly_curve: bool = False
    center_line: bool = False


def draw_overlay(
    frame: np.ndarray,
    result: LineDetectResult | None,
    flags: OverlayFlags,
) -> np.ndarray:
    """カメラ映像にオーバーレイを描画する

    Args:
        frame: 元の BGR カメラ画像
        result: 線検出の結果(None の場合はオーバーレイなし)
        flags: 表示項目のフラグ

    Returns:
        オーバーレイ描画済みの画像
    """
    display = frame.copy()
    h, w = display.shape[:2]

    if result is None:
        return display

    # 二値化画像の半透明オーバーレイ
    if flags.binary and result.binary_image is not None:
        display = _draw_binary_overlay(
            display, result.binary_image,
        )

    # 検出領域の枠
    if flags.detect_region:
        cv2.rectangle(
            display,
            (0, 0),
            (w - 1, h - 1),
            COLOR_REGION, 1,
        )

    # 画像中心線
    if flags.center_line:
        center_x = w // 2
        cv2.line(
            display,
            (center_x, 0),
            (center_x, h),
            COLOR_CENTER, 1,
        )

    # フィッティング曲線
    if flags.poly_curve and result.poly_coeffs is not None:
        _draw_poly_curve(display, result.poly_coeffs)

    return display


def _draw_binary_overlay(
    frame: np.ndarray,
    binary: np.ndarray,
) -> np.ndarray:
    """二値化画像を半透明で重ねる

    Args:
        frame: 元の BGR 画像
        binary: 二値化画像(グレースケール)

    Returns:
        合成された画像
    """
    binary_bgr = np.zeros_like(frame)
    binary_bgr[:, :, 2] = binary

    return cv2.addWeighted(
        frame, 1.0 - BINARY_OPACITY,
        binary_bgr, BINARY_OPACITY, 0,
    )


def _draw_poly_curve(
    frame: np.ndarray,
    coeffs: np.ndarray,
) -> None:
    """フィッティング曲線を描画する

    Args:
        frame: 描画先の画像
        coeffs: 多項式の係数
    """
    h, w = frame.shape[:2]
    poly = np.poly1d(coeffs)

    # 曲線上の点を生成
    ys = np.arange(0, h)
    xs = poly(ys)

    # 画像範囲内の点のみ描画
    points = []
    for x, y in zip(xs, ys):
        ix = int(round(x))
        if 0 <= ix < w:
            points.append([ix, int(y)])

    if len(points) >= 2:
        pts = np.array(points, dtype=np.int32)
        cv2.polylines(
            frame, [pts], False,
            COLOR_LINE, 1,
        )