app/core/utils.py - master - ksrk/layout_validation

app/core/utils.py
266 строк11 KB
ksrk
refactoring: introduced docstrings in some classes and complicated utils
03 июл 2026, 00:03
03 июл 2026, 00:035367527
from typing import Tuple

import cv2
import fitz
import numpy as np
from PIL import Image, ImageCms
from PIL.ImageFile import ImageFile
from pymupdf import Page


def extract_contrast_color(img: np.ndarray | cv2.Mat) -> Tuple[int, int, int]:
    """Рассчитывает контрастный изображению цвет выделения,
    используя цветовое представление HSV (Hue, Saturation, Value).

    :param img: Исходное изображение (матрица OpenCV/NumPy), для которого
        необходимо подобрать контрастный цвет выделения.
    :return: Кортеж из трех целых чисел, представляющий контрастный цвет в формате BGR
        (Blue, Green, Red), готовый для использования в функциях рисования OpenCV.
    """

    img_resized = cv2.resize(img, (64, 64), interpolation=cv2.INTER_AREA)
    hsv = cv2.cvtColor(img_resized, cv2.COLOR_BGR2HSV)

    h_channel = hsv[:, :, 0].flatten()
    s_channel = hsv[:, :, 1].flatten()
    colored_pixels = h_channel[s_channel > 30]
    if len(colored_pixels) == 0:
        avg_v = np.mean(hsv[:, :, 2])
        contrast_h = 30 if avg_v < 128 else 120
        contrast_s, contrast_v = 255, 255
    else:
        hist, _ = np.histogram(colored_pixels, bins=180, range=(0, 180))
        hist_smooth = cv2.GaussianBlur(hist.astype(np.float32), (1, 15), 0).flatten()
        contrast_h = int(np.argmin(hist_smooth))
        contrast_s = 255
        contrast_v = 255

    contrast_hsv_pixel = np.uint8([[[contrast_h, contrast_s, contrast_v]]])
    contrast_bgr_pixel = cv2.cvtColor(contrast_hsv_pixel, cv2.COLOR_HSV2BGR)
    b, g, r = contrast_bgr_pixel.ravel()

    return int(b), int(g), int(r)


def draw_violation_arrow(
    canvas: np.ndarray,
    violation_type: str,
    rect_px: Tuple[Tuple[int, int], Tuple[int, int]],
    trim_box_px: Tuple[Tuple[int, int], Tuple[int, int]],
    color: Tuple[int, int, int] = (0, 255, 0),
    dpi: int = 300
):
    """Рисует перпендикулярную стрелку от ближайшей точки элемента
    до ближайшей стороны края вырубки.

    :param canvas: Исходное изображение (массив OpenCV/NumPy), на котором будет выполнено рисование.
    :param violation_type: Тип нарушения, определяющий направление стрелки. Если значение
        равно "outer", стрелка рисуется от края вырубки к элементу, иначе — от элемента к краю.
    :param rect_px: Координаты ограничивающего прямоугольника (bounding box) проблемного элемента в пикселях.
        Принимает формат пары точек ((x0, y0), (x1, y1)) или плоский список [x0, y0, x1, y1].
    :param trim_box_px: Координаты границ обрезного формата (TrimBox / края вырубки) в пикселях.
        Принимает формат пары точек ((tx0, ty0), (tx1, ty1)) или плоский список [tx0, ty0, tx1, ty1].
    :param color: Цвет в формате BGR для отрисовки прямоугольника вокруг элемента.
        По умолчанию (0, 255, 0) (зеленый).
    :param dpi: Разрешение изображения в DPI (точек на дюйм) для корректного перевода
        расстояния из пикселей в миллиметры. По умолчанию 300.
    """

    if isinstance(rect_px, (list, tuple)) and isinstance(rect_px[0], (list, tuple)):
        (ex0, ey0), (ex1, ey1) = rect_px
    else:
        ex0, ey0, ex1, ey1 = rect_px

    if isinstance(trim_box_px, (list, tuple)) and isinstance(trim_box_px[0], (list, tuple)):
        (tx0, ty0), (tx1, ty1) = trim_box_px
    else:
        tx0, ty0, tx1, ty1 = trim_box_px
    ex0, ey0, ex1, ey1 = int(ex0), int(ey0), int(ex1), int(ey1)
    tx0, ty0, tx1, ty1 = int(tx0), int(ty0), int(tx1), int(ty1)

    dist_top = abs(ey0 - ty0)
    dist_bottom = abs(ey1 - ty1)
    dist_left = abs(ex0 - tx0)
    dist_right = abs(ex1 - tx1)

    distances = {
        "top": dist_top,
        "bottom": dist_bottom,
        "left": dist_left,
        "right": dist_right
    }

    side = min(distances, key=distances.get)

    if side == "top":
        object_edge_center = ((ex0 + ex1) // 2, ey0)
        trim_edge_pt = (object_edge_center[0], ty0)

    elif side == "bottom":
        object_edge_center = ((ex0 + ex1) // 2, ey1)
        trim_edge_pt = (object_edge_center[0], ty1)

    elif side == "left":
        object_edge_center = (ex0, (ey0 + ey1) // 2)
        trim_edge_pt = (tx0, object_edge_center[1])

    else:
        object_edge_center = (ex1, (ey0 + ey1) // 2)
        trim_edge_pt = (tx1, object_edge_center[1])

    if violation_type == "outer":
        pt_start = trim_edge_pt
        pt_end = object_edge_center
    else:
        pt_start = object_edge_center
        pt_end = trim_edge_pt

    cv2.rectangle(canvas, (ex0, ey0), (ex1, ey1), color, 2)

    cv2.arrowedLine(
        canvas,
        pt_start,
        pt_end,
        (0, 0, 0),
        thickness=2,
        tipLength=0.2
    )
    mid_x = (pt_start[0] + pt_end[0]) // 2
    mid_y = (pt_start[1] + pt_end[1]) // 2

    dist_px = int(np.sqrt((pt_end[0] - pt_start[0]) ** 2 + (pt_end[1] - pt_start[1]) ** 2))
    dist_mm = dist_px * 25.4 / dpi

    text_offset_x = 8 if side in ("top", "bottom") else 0
    text_offset_y = -5 if side in ("left", "right") else -5

    cv2.putText(
        canvas,
        f"{dist_mm:.2f}mm",
        (mid_x + text_offset_x, mid_y + text_offset_y),
        cv2.FONT_HERSHEY_SIMPLEX,
        1.0,
        (0, 0, 0),
        2,
        cv2.LINE_AA
    )


def rgb_to_cmyk_gcr(page: Page, dpi: int=300) -> Image.Image:
    """Переводит цвета на векторном изображении из RGB в CMYK.

    :param page: Объект страницы документа для растеризации и цветоделения.
    :param dpi: Разрешение рендеринга страницы в точках на дюйм (по умолчанию 300).
    :return: Растровое изображение страницы в цветовом пространстве CMYK.
    """

    mat = fitz.Matrix(dpi / 72, dpi / 72)
    pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB, alpha=False)

    img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
    rgb = np.array(img, dtype=np.float32) / 255.0
    c_temp = 1.0 - rgb[..., 0]
    m_temp = 1.0 - rgb[..., 1]
    y_temp = 1.0 - rgb[..., 2]

    k = np.minimum(np.minimum(c_temp, m_temp), y_temp)
    pure_black = (k == 1.0)
    denominator = np.where(pure_black, 1.0, 1.0 - k)

    c = np.where(pure_black, 0.0, (c_temp - k) / denominator)
    m = np.where(pure_black, 0.0, (m_temp - k) / denominator)
    y = np.where(pure_black, 0.0, (y_temp - k) / denominator)
    cmyk = (np.stack([c, m, y, k], axis=-1) * 255.0).astype(np.uint8)
    return Image.fromarray(cmyk, mode="CMYK")


def convert_rgb_to_cmyk(page: Page, dpi: int = 300) -> Image.Image:
    """Переводит цвета на векторном изображении из RGB в CMYK."""

    mat = fitz.Matrix(dpi / 72, dpi / 72)
    pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB, alpha=False)
    img = Image.frombytes(
        "RGB",
        (pix.width, pix.height),
        pix.samples
    )
    srgb_profile = ImageCms.createProfile('sRGB')
    cmyk_profile = ImageCms.getOpenProfile('./profiles/ISOcoated_v2_300_eci.icc')
    img_cmyk = ImageCms.profileToProfile(
        img,
        srgb_profile,
        cmyk_profile,
        outputMode='CMYK'
    )
    return img_cmyk


def compute_img_size(img: ImageFile) -> Tuple[int, int]:
    """Вычисляет физический размер изображения в миллиметрах."""

    dpi = img.info.get('dpi', (300, 300)) # если параметр dpi не зашит, то тогда пока что принимаем по дефолту 300
    width_mm = (img.width * 25.4) / dpi[0]
    height_mm = (img.height * 25.4) / dpi[1]
    return int(width_mm), int(height_mm)


def compute_pix_size(img: ImageFile, dpi: int = 300) -> Tuple[float, float]:
    """Вычисляет физический размер одного пикселя в миллиметрах."""

    dpi = img.info.get('dpi', (dpi, dpi))
    pix_w = 25.4 / dpi[0]
    pix_h = 25.4 / dpi[1]
    return pix_w, pix_h


def convert_to_cmyk_raster(page: Page, dpi: int = 300) -> Image.Image:
    mat = fitz.Matrix(dpi / 72, dpi / 72)
    pix = page.get_pixmap(matrix=mat, colorspace=fitz.csCMYK, alpha=False)
    img = Image.frombytes(
        "CMYK",
        (pix.width, pix.height),
        pix.samples
    )
    return img


def as_opencv_img(page: Page, dpi=300) -> cv2.Mat | np.ndarray:
    """Конвертирует страницу векторного макета на
    растровое представление для работы с OpenCV.

    :param page: Объект страницы документа для растеризации.
    :param dpi: Разрешение рендеринга страницы в точках на дюйм (по умолчанию 300).
    :return: Изображение страницы в формате матрицы OpenCV (обычно в цветовом пространстве BGR).
    """

    mat = fitz.Matrix(dpi / 72, dpi / 72)
    pix = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB, alpha=False)
    img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
    return img


def from_pt_to_mm(value_pt) -> float:
    """Переводит величину из пунктов (1/72 дюйма) в миллиметры"""

    return value_pt * 25.4 / 72


def from_mm_to_pt(value_mm: float) -> float:
    """Переводит величину из миллиметров в пункты (1/72 дюйма)"""

    return value_mm * 72 / 25.4


def from_pt_to_px(value_px: float) -> int:
    """Переводит величину из пунктов (1/72 дюйма) в пиксели"""

    zoom = 300 / 72
    return int(value_px * zoom)


def from_px_to_pt(value_px: float) -> int:
    """Переводит величину из пикселей в пункты (1/72 дюйма)"""

    zoom = 72 / 300
    return int(value_px * zoom)