mediaproc/gpu/models/cv/edges.py

"""
Edge detection — Canny + HoughLinesP → parallel line pairs → bounding boxes.

Finds horizontal line pairs with consistent spacing, which correspond to
the top and bottom edges of advertising hoardings.
"""

from __future__ import annotations

import base64
import io

import cv2
import numpy as np


def detect_edges(
    image: np.ndarray,
    canny_low: int = 50,
    canny_high: int = 150,
    hough_threshold: int = 80,
    hough_min_length: int = 100,
    hough_max_gap: int = 10,
    pair_max_distance: int = 200,
    pair_min_distance: int = 15,
) -> list[dict]:
    """
    Find horizontal line pairs that likely bound advertising hoardings.

    Returns list of dicts with keys: x, y, w, h, confidence, label.
    Each box represents the region between a detected pair of parallel
    horizontal lines.
    """
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, canny_low, canny_high)

    raw_lines = cv2.HoughLinesP(
        edges,
        rho=1,
        theta=np.pi / 180,
        threshold=hough_threshold,
        minLineLength=hough_min_length,
        maxLineGap=hough_max_gap,
    )

    if raw_lines is None:
        return []

    # Filter to near-horizontal lines (within 10 degrees)
    horizontals = _filter_horizontal(raw_lines, max_angle_deg=10)

    if len(horizontals) < 2:
        return []

    # Find pairs of parallel horizontals with consistent spacing
    pairs = _find_line_pairs(
        horizontals,
        min_distance=pair_min_distance,
        max_distance=pair_max_distance,
    )

    # Convert pairs to bounding boxes
    h, w = image.shape[:2]
    results = []
    for top_line, bottom_line in pairs:
        box = _pair_to_bbox(top_line, bottom_line, frame_width=w, frame_height=h)
        if box is not None:
            results.append(box)

    return results


def _filter_horizontal(lines: np.ndarray, max_angle_deg: float = 10) -> list[tuple]:
    """Keep only lines within max_angle_deg of horizontal."""
    max_slope = np.tan(np.radians(max_angle_deg))
    result = []
    for line in lines:
        x1, y1, x2, y2 = line[0]
        dx = x2 - x1
        if dx == 0:
            continue
        slope = abs((y2 - y1) / dx)
        if slope <= max_slope:
            y_mid = (y1 + y2) / 2
            x_min = min(x1, x2)
            x_max = max(x1, x2)
            length = np.sqrt(dx**2 + (y2 - y1) ** 2)
            result.append((x_min, x_max, y_mid, length))
    return result


def _find_line_pairs(
    horizontals: list[tuple],
    min_distance: int,
    max_distance: int,
) -> list[tuple]:
    """
    Find pairs of horizontal lines that could be top/bottom of a hoarding.

    Lines must overlap horizontally and be spaced within [min_distance, max_distance].
    """
    # Sort by y position
    sorted_lines = sorted(horizontals, key=lambda l: l[2])

    pairs = []
    used = set()

    for i, top in enumerate(sorted_lines):
        if i in used:
            continue
        for j, bottom in enumerate(sorted_lines[i + 1 :], start=i + 1):
            if j in used:
                continue

            y_gap = bottom[2] - top[2]
            if y_gap < min_distance:
                continue
            if y_gap > max_distance:
                break  # sorted by y, no point checking further

            # Check horizontal overlap
            overlap_start = max(top[0], bottom[0])
            overlap_end = min(top[1], bottom[1])
            overlap = overlap_end - overlap_start

            # Require at least 50% overlap relative to shorter line
            shorter_length = min(top[1] - top[0], bottom[1] - bottom[0])
            if shorter_length > 0 and overlap / shorter_length >= 0.5:
                pairs.append((top, bottom))
                used.add(i)
                used.add(j)
                break

    return pairs


def _pair_to_bbox(
    top: tuple,
    bottom: tuple,
    frame_width: int,
    frame_height: int,
) -> dict | None:
    """Convert a line pair to a bounding box dict."""
    x = int(max(0, min(top[0], bottom[0])))
    y = int(max(0, top[2]))
    x2 = int(min(frame_width, max(top[1], bottom[1])))
    y2 = int(min(frame_height, bottom[2]))
    w = x2 - x
    h = y2 - y

    if w < 20 or h < 5:
        return None

    # Confidence based on line lengths relative to box width
    avg_line_length = (top[3] + bottom[3]) / 2
    coverage = min(1.0, avg_line_length / max(w, 1))

    return {
        "x": x,
        "y": y,
        "w": w,
        "h": h,
        "confidence": round(coverage, 3),
        "label": "edge_region",
    }


def _np_to_b64_jpeg(image: np.ndarray, quality: int = 70) -> str:
    """Encode a numpy image (BGR or grayscale) as base64 JPEG."""
    ok, buf = cv2.imencode(".jpg", image, [cv2.IMWRITE_JPEG_QUALITY, quality])
    if not ok:
        return ""
    return base64.b64encode(buf.tobytes()).decode()


def detect_edges_debug(
    image: np.ndarray,
    canny_low: int = 50,
    canny_high: int = 150,
    hough_threshold: int = 80,
    hough_min_length: int = 100,
    hough_max_gap: int = 10,
    pair_max_distance: int = 200,
    pair_min_distance: int = 15,
) -> dict:
    """
    Same as detect_edges but returns intermediate visualizations.

    Returns dict with:
        regions: list[dict]          — same boxes as detect_edges
        edge_overlay_b64: str        — Canny edge image as base64 JPEG
        lines_overlay_b64: str       — frame with Hough lines drawn
        horizontal_count: int        — number of horizontal lines found
        pair_count: int              — number of line pairs found
    """
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, canny_low, canny_high)

    # Edge overlay — Canny output as-is (white edges on black)
    edge_overlay_b64 = _np_to_b64_jpeg(edges)

    raw_lines = cv2.HoughLinesP(
        edges,
        rho=1,
        theta=np.pi / 180,
        threshold=hough_threshold,
        minLineLength=hough_min_length,
        maxLineGap=hough_max_gap,
    )

    # Lines overlay — draw all Hough lines on a copy of the frame
    lines_vis = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    if raw_lines is not None:
        for line in raw_lines:
            x1, y1, x2, y2 = line[0]
            cv2.line(lines_vis, (x1, y1), (x2, y2), (0, 0, 255), 1)

    horizontals = []
    if raw_lines is not None:
        horizontals = _filter_horizontal(raw_lines, max_angle_deg=10)

    # Draw horizontal lines in cyan, thicker
    for h_line in horizontals:
        x_min, x_max, y_mid, _ = h_line
        cv2.line(lines_vis, (int(x_min), int(y_mid)), (int(x_max), int(y_mid)), (255, 255, 0), 2)

    pairs = []
    if len(horizontals) >= 2:
        pairs = _find_line_pairs(
            horizontals,
            min_distance=pair_min_distance,
            max_distance=pair_max_distance,
        )

    # Draw paired lines in green
    for top_line, bottom_line in pairs:
        cv2.line(lines_vis, (int(top_line[0]), int(top_line[2])),
                 (int(top_line[1]), int(top_line[2])), (0, 255, 0), 2)
        cv2.line(lines_vis, (int(bottom_line[0]), int(bottom_line[2])),
                 (int(bottom_line[1]), int(bottom_line[2])), (0, 255, 0), 2)

    lines_overlay_b64 = _np_to_b64_jpeg(lines_vis)

    # Build region boxes (same logic as detect_edges)
    h, w = image.shape[:2]
    regions = []
    for top_line, bottom_line in pairs:
        box = _pair_to_bbox(top_line, bottom_line, frame_width=w, frame_height=h)
        if box is not None:
            regions.append(box)

    return {
        "regions": regions,
        "edge_overlay_b64": edge_overlay_b64,
        "lines_overlay_b64": lines_overlay_b64,
        "horizontal_count": len(horizontals),
        "pair_count": len(pairs),
    }