cv_ops/transforms.py

"""2D geometric transformations implemented with OpenCV."""

from __future__ import annotations

import cv2
import numpy as np

from filters.builtin import ensure_rgb


BORDER_MODES = {"constant": cv2.BORDER_CONSTANT, "reflect": cv2.BORDER_REFLECT, "replicate": cv2.BORDER_REPLICATE}
INTERPOLATION = {"nearest": cv2.INTER_NEAREST, "linear": cv2.INTER_LINEAR, "cubic": cv2.INTER_CUBIC, "area": cv2.INTER_AREA}


def translate(image: np.ndarray, x: int = 20, y: int = 20, border: str = "constant") -> tuple[np.ndarray, list[list[float]]]:
    img = ensure_rgb(image)
    matrix = np.float32([[1, 0, x], [0, 1, y]])
    out = cv2.warpAffine(img, matrix, (img.shape[1], img.shape[0]), borderMode=BORDER_MODES.get(border, cv2.BORDER_CONSTANT))
    return out, matrix.tolist()


def rotate(image: np.ndarray, angle: float = 30, scale: float = 1.0, center_x: float = 0.5, center_y: float = 0.5, expand: bool = True) -> tuple[np.ndarray, list[list[float]]]:
    img = ensure_rgb(image)
    h, w = img.shape[:2]
    center = (w * center_x, h * center_y)
    matrix = cv2.getRotationMatrix2D(center, angle, scale)
    out_w, out_h = w, h
    if expand:
        cos, sin = abs(matrix[0, 0]), abs(matrix[0, 1])
        out_w, out_h = int((h * sin) + (w * cos)), int((h * cos) + (w * sin))
        matrix[0, 2] += out_w / 2 - center[0]
        matrix[1, 2] += out_h / 2 - center[1]
    return cv2.warpAffine(img, matrix, (out_w, out_h)), matrix.tolist()


def scale_image(image: np.ndarray, sx: float = 1.2, sy: float = 1.2, interpolation: str = "linear") -> tuple[np.ndarray, list[list[float]]]:
    img = ensure_rgb(image)
    out = cv2.resize(img, None, fx=sx, fy=sy, interpolation=INTERPOLATION.get(interpolation, cv2.INTER_LINEAR))
    return out, [[sx, 0, 0], [0, sy, 0]]


def reflect(image: np.ndarray, mode: str = "horizontal") -> tuple[np.ndarray, list[list[float]]]:
    img = ensure_rgb(image)
    code = {"horizontal": 1, "vertical": 0, "both": -1}.get(mode, 1)
    matrix = {"horizontal": [[-1, 0, img.shape[1]], [0, 1, 0]], "vertical": [[1, 0, 0], [0, -1, img.shape[0]]], "both": [[-1, 0, img.shape[1]], [0, -1, img.shape[0]]]}.get(mode)
    return cv2.flip(img, code), matrix