grid2frame/extractor.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

import cv2
import numpy as np


@dataclass(frozen=True)
class FrameRegion:
    id: str
    row: int
    col: int
    x: int
    y: int
    width: int
    height: int


@dataclass(frozen=True)
class DetectionResult:
    regions: list[FrameRegion]
    rows: int
    columns: int
    width: int
    height: int


@dataclass
class Band:
    start: int
    end: int


@dataclass(frozen=True)
class Segment:
    start: int
    end: int


def read_image(path: str | Path) -> np.ndarray:
    image = cv2.imdecode(np.fromfile(Path(path), dtype=np.uint8), cv2.IMREAD_COLOR)
    if image is None:
        raise ValueError(f"Could not read image: {path}")
    return image


def detect_grid(
    image: np.ndarray,
    min_frame_size: int = 80,
    sensitivity: int = 58,
    separator_mode: str = "auto",
) -> DetectionResult:
    height, width = image.shape[:2]
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    vertical_bands = find_separator_bands(
        gray=gray,
        axis="vertical",
        min_frame_size=min_frame_size,
        sensitivity=sensitivity,
        separator_mode=separator_mode,
    )
    horizontal_bands = find_separator_bands(
        gray=gray,
        axis="horizontal",
        min_frame_size=min_frame_size,
        sensitivity=sensitivity,
        separator_mode=separator_mode,
    )

    x_segments = bands_to_segments(vertical_bands, width, min_frame_size)
    y_segments = bands_to_segments(horizontal_bands, height, min_frame_size)
    regions: list[FrameRegion] = []

    for row, y_segment in enumerate(y_segments):
        for col, x_segment in enumerate(x_segments):
            regions.append(
                FrameRegion(
                    id=f"{row + 1}-{col + 1}",
                    row=row,
                    col=col,
                    x=x_segment.start,
                    y=y_segment.start,
                    width=x_segment.end - x_segment.start,
                    height=y_segment.end - y_segment.start,
                )
            )

    return DetectionResult(
        regions=regions,
        rows=len(y_segments),
        columns=len(x_segments),
        width=width,
        height=height,
    )


def extract_frames(image: np.ndarray, regions: list[FrameRegion]) -> list[np.ndarray]:
    frames = []
    for region in regions:
        frames.append(
            image[
                region.y : region.y + region.height,
                region.x : region.x + region.width,
            ].copy()
        )
    return frames


def annotate_regions(image: np.ndarray, regions: list[FrameRegion]) -> np.ndarray:
    annotated = image.copy()
    color = (190, 255, 55)

    for region in regions:
        cv2.rectangle(
            annotated,
            (region.x, region.y),
            (region.x + region.width, region.y + region.height),
            color,
            max(1, round(max(image.shape[:2]) / 700)),
        )

    return annotated


def find_separator_bands(
    gray: np.ndarray,
    axis: str,
    min_frame_size: int,
    sensitivity: int,
    separator_mode: str,
) -> list[Band]:
    sensitivity = max(0, min(100, sensitivity))
    dark_threshold = 10 + round(sensitivity * 0.24)
    light_threshold = 246 - round(sensitivity * 0.16)
    line_ratio_threshold = 0.972 - sensitivity * 0.00035

    if separator_mode == "dark":
        separator_mask = gray <= dark_threshold
    elif separator_mode == "light":
        separator_mask = gray >= light_threshold
    else:
        separator_mask = (gray <= dark_threshold) | (gray >= light_threshold)

    if axis == "vertical":
        profile = separator_mask.mean(axis=0)
        axis_size = gray.shape[1]
    else:
        profile = separator_mask.mean(axis=1)
        axis_size = gray.shape[0]

    candidate_indexes = np.flatnonzero(profile >= line_ratio_threshold)
    bands = group_indexes(candidate_indexes)
    max_thickness = max(2, min(round(axis_size * 0.018), round(min_frame_size * 0.35)))
    bands = [band for band in bands if band.end - band.start + 1 <= max_thickness]

    return merge_close_bands(bands, min_gap=max(2, round(min_frame_size * 0.08)))


def group_indexes(indexes: np.ndarray) -> list[Band]:
    if indexes.size == 0:
        return []

    bands: list[Band] = []
    start = int(indexes[0])
    previous = int(indexes[0])

    for raw_index in indexes[1:]:
        index = int(raw_index)
        if index <= previous + 2:
            previous = index
            continue

        bands.append(Band(start=start, end=previous))
        start = index
        previous = index

    bands.append(Band(start=start, end=previous))
    return bands


def merge_close_bands(bands: list[Band], min_gap: int) -> list[Band]:
    if not bands:
        return []

    merged = [bands[0]]
    for band in bands[1:]:
        current = merged[-1]
        if band.start - current.end <= min_gap:
            current.end = band.end
        else:
            merged.append(band)

    return merged


def bands_to_segments(bands: list[Band], axis_size: int, min_frame_size: int) -> list[Segment]:
    if not bands:
        return [Segment(start=0, end=axis_size)]

    edge_tolerance = max(3, round(axis_size * 0.008))
    first_pixel = 0
    last_pixel = axis_size
    internal_bands: list[Band] = []

    for band in bands:
        if band.start <= edge_tolerance:
            first_pixel = max(first_pixel, band.end + 1)
        elif band.end >= axis_size - edge_tolerance - 1:
            last_pixel = min(last_pixel, band.start)
        else:
            internal_bands.append(band)

    segments: list[Segment] = []
    segment_start = first_pixel

    for band in internal_bands:
        push_segment(segments, segment_start, band.start, min_frame_size)
        segment_start = band.end + 1

    push_segment(segments, segment_start, last_pixel, min_frame_size)
    return segments


def push_segment(segments: list[Segment], start: int, end: int, min_frame_size: int) -> None:
    start = max(0, int(round(start)))
    end = max(start, int(round(end)))

    if end - start >= min_frame_size:
        segments.append(Segment(start=start, end=end))