src/photo_processor/segmentation.py - master - sarus/photo-processor

src/photo_processor/segmentation.py
494 строки17 KB
otcheskiy
feat: add simple background segmentation and mask quality
01 авг 2026, 12:34
01 авг 2026, 12:3403c5abe
"""Построение и очистка простой бинарной маски объекта."""

from __future__ import annotations

from dataclasses import dataclass

import cv2
import numpy as np
from PIL import Image

from photo_processor.background import BackgroundStats, _luminance
from photo_processor.config import MaskConfig


@dataclass(frozen=True)
class ComponentInfo:
    label: int
    area: int
    x: int
    y: int
    width: int
    height: int
    touches_border: bool

    @property
    def aspect_ratio(self) -> float:
        short = max(1, min(self.width, self.height))
        return max(self.width, self.height) / float(short)


@dataclass(frozen=True)
class SegmentationResult:
    raw_mask: np.ndarray
    clean_mask: np.ndarray
    raw_foreground_ratio: float
    clean_foreground_ratio: float
    component_count: int
    significant_component_count: int
    border_component_count: int
    removed_border_component_count: int
    removed_component_count: int
    removed_foreground_pixels: int
    foreground_retention_ratio: float
    foreground_loss_ratio: float
    object_extent_retention_ratio: float
    large_border_touch_kept: bool
    components: list[ComponentInfo]


def _rgb_to_lab(rgb: np.ndarray) -> np.ndarray:
    """RGB uint8/float → LAB float32 через OpenCV."""
    if rgb.dtype != np.uint8:
        rgb_u8 = np.clip(rgb, 0, 255).astype(np.uint8)
    else:
        rgb_u8 = rgb
    bgr = cv2.cvtColor(rgb_u8, cv2.COLOR_RGB2BGR)
    lab = cv2.cvtColor(bgr, cv2.COLOR_BGR2LAB).astype(np.float32)
    return lab


def _mask_bbox(mask: np.ndarray) -> tuple[int, int, int, int] | None:
    ys, xs = np.where(mask > 0)
    if ys.size == 0:
        return None
    x0, x1 = int(xs.min()), int(xs.max())
    y0, y1 = int(ys.min()), int(ys.max())
    return x0, y0, x1 - x0 + 1, y1 - y0 + 1


def build_raw_mask(
    image: Image.Image,
    background: BackgroundStats,
    config: MaskConfig,
) -> np.ndarray:
    """Сырая маска: 0 = фон, 255 = объект."""
    rgb = np.asarray(image.convert("RGB"), dtype=np.float32)
    bg = np.array(background.median_rgb, dtype=np.float32)
    color_dist = np.linalg.norm(rgb - bg, axis=2)
    brightness = _luminance(rgb)

    lab = _rgb_to_lab(rgb)
    bg_lab = _rgb_to_lab(bg.reshape(1, 1, 3))[0, 0]
    lab_dist = np.linalg.norm(lab - bg_lab, axis=2)

    dark_cut = min(
        config.darkness_threshold,
        background.brightness - config.relative_darkness_margin,
    )
    foreground = (
        (color_dist >= config.color_distance_threshold)
        | (brightness <= dark_cut)
        | (lab_dist >= config.lab_distance_threshold)
    )
    mask = np.where(foreground, 255, 0).astype(np.uint8)

    # Итеративное дотягивание тонких линий от уже найденного foreground.
    dilate_px = max(0, config.thin_reclaim_dilate_px)
    iterations = max(0, config.thin_reclaim_iterations)
    if dilate_px > 0 and iterations > 0:
        kernel = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (dilate_px * 2 + 1, dilate_px * 2 + 1)
        )
        for _ in range(iterations):
            vicinity = cv2.dilate(mask, kernel, iterations=1)
            reclaim = (
                (vicinity > 0)
                & (mask == 0)
                & (
                    (color_dist >= config.thin_reclaim_color_distance)
                    | (brightness <= dark_cut)
                    | (lab_dist >= config.thin_reclaim_lab_distance)
                )
            )
            if not np.any(reclaim):
                break
            mask = np.where(reclaim | (mask > 0), 255, 0).astype(np.uint8)

    return mask


def build_structure_probe_mask(
    image: Image.Image,
    background: BackgroundStats,
    config: MaskConfig,
) -> np.ndarray:
    """Более чувствительная маска только для оценки полноты тонких структур."""
    rgb = np.asarray(image.convert("RGB"), dtype=np.float32)
    bg = np.array(background.median_rgb, dtype=np.float32)
    color_dist = np.linalg.norm(rgb - bg, axis=2)
    brightness = _luminance(rgb)
    lab = _rgb_to_lab(rgb)
    bg_lab = _rgb_to_lab(bg.reshape(1, 1, 3))[0, 0]
    lab_dist = np.linalg.norm(lab - bg_lab, axis=2)
    dark_cut = min(
        config.darkness_threshold,
        background.brightness - config.structure_probe_relative_darkness_margin,
    )
    foreground = (
        (color_dist >= config.structure_probe_color_distance)
        | (brightness <= dark_cut)
        | (lab_dist >= config.structure_probe_lab_distance)
    )
    return np.where(foreground, 255, 0).astype(np.uint8)


def _analyze_components(mask: np.ndarray) -> tuple[np.ndarray, list[ComponentInfo]]:
    num_labels, labels, stats, _centroids = cv2.connectedComponentsWithStats(
        mask, connectivity=8
    )
    height, width = mask.shape
    components: list[ComponentInfo] = []
    for label in range(1, num_labels):
        area = int(stats[label, cv2.CC_STAT_AREA])
        x = int(stats[label, cv2.CC_STAT_LEFT])
        y = int(stats[label, cv2.CC_STAT_TOP])
        w = int(stats[label, cv2.CC_STAT_WIDTH])
        h = int(stats[label, cv2.CC_STAT_HEIGHT])
        touches = (
            x <= 0
            or y <= 0
            or (x + w) >= width
            or (y + h) >= height
            or bool(np.any(labels[0, :] == label))
            or bool(np.any(labels[-1, :] == label))
            or bool(np.any(labels[:, 0] == label))
            or bool(np.any(labels[:, -1] == label))
        )
        components.append(
            ComponentInfo(
                label=label,
                area=area,
                x=x,
                y=y,
                width=w,
                height=h,
                touches_border=touches,
            )
        )
    return labels, components


def _expanded_bbox(
    bbox: tuple[int, int, int, int],
    expand: int,
    width: int,
    height: int,
) -> tuple[int, int, int, int]:
    x, y, w, h = bbox
    x0 = max(0, x - expand)
    y0 = max(0, y - expand)
    x1 = min(width - 1, x + w - 1 + expand)
    y1 = min(height - 1, y + h - 1 + expand)
    return x0, y0, x1 - x0 + 1, y1 - y0 + 1


def _component_near_or_inside(
    component: ComponentInfo,
    bbox: tuple[int, int, int, int],
) -> bool:
    bx, by, bw, bh = bbox
    cx0, cy0 = component.x, component.y
    cx1, cy1 = component.x + component.width - 1, component.y + component.height - 1
    ix0 = max(bx, cx0)
    iy0 = max(by, cy0)
    ix1 = min(bx + bw - 1, cx1)
    iy1 = min(by + bh - 1, cy1)
    return ix0 <= ix1 and iy0 <= iy1


def _component_center_inside(
    component: ComponentInfo,
    bbox: tuple[int, int, int, int],
) -> bool:
    bx, by, bw, bh = bbox
    cx = component.x + component.width / 2.0
    cy = component.y + component.height / 2.0
    return bx <= cx <= bx + bw - 1 and by <= cy <= by + bh - 1


def _is_elongated(component: ComponentInfo, config: MaskConfig) -> bool:
    return component.aspect_ratio >= config.min_elongation_aspect_ratio


def _should_keep_detail(
    component: ComponentInfo,
    *,
    expanded: tuple[int, int, int, int],
    near_expanded: tuple[int, int, int, int],
    config: MaskConfig,
) -> bool:
    elongated = _is_elongated(component, config)

    if component.area >= config.min_component_area and _component_near_or_inside(
        component, near_expanded
    ):
        return True

    # Компактные внутренние детали (логотип) — только в плотной зоне оправы.
    if (
        not elongated
        and component.area >= config.min_internal_component_area
        and _component_near_or_inside(component, expanded)
    ):
        return True

    # Тонкие/вытянутые фрагменты дужек — допускаем более широкий near-bbox.
    if (
        elongated
        and component.area >= config.min_elongated_component_area
        and _component_near_or_inside(component, near_expanded)
    ):
        return True
    return False


def _non_border_mask(mask: np.ndarray) -> np.ndarray:
    labels, components = _analyze_components(mask)
    keep = {c.label for c in components if not c.touches_border}
    result = np.zeros_like(mask)
    if keep:
        result[np.isin(labels, list(keep))] = 255
    return result


def clean_mask(
    raw_mask: np.ndarray,
    config: MaskConfig,
    *,
    structure_reference: np.ndarray | None = None,
) -> SegmentationResult:
    """Очистить маску: удалить краевые артефакты, шум, сохранить детали оправы."""
    height, width = raw_mask.shape
    image_area = float(height * width)
    raw_ratio = float(np.count_nonzero(raw_mask) / image_area) if image_area else 0.0

    labels, components = _analyze_components(raw_mask)
    border_component_count = sum(1 for c in components if c.touches_border)
    initial_component_count = len(components)

    def _center_of_mass_in_central_zone(component: ComponentInfo) -> bool:
        cx = component.x + component.width / 2.0
        cy = component.y + component.height / 2.0
        return (
            width * 0.15 <= cx <= width * 0.85
            and height * 0.15 <= cy <= height * 0.85
        )

    object_candidates = [
        c
        for c in components
        if c.area >= config.min_component_area and not c.touches_border
    ]
    if not object_candidates:
        object_candidates = [
            c
            for c in components
            if c.area >= config.min_component_area
            and _center_of_mass_in_central_zone(c)
        ]
    if not object_candidates:
        object_candidates = [
            c for c in components if c.area >= config.min_component_area
        ]

    keep: set[int] = set()
    removed_border = 0
    large_border_kept = False

    if object_candidates:
        main = max(object_candidates, key=lambda item: item.area)
        keep.add(main.label)
        if main.touches_border:
            large_border_kept = True

        main_bbox = (main.x, main.y, main.width, main.height)
        expanded = _expanded_bbox(
            main_bbox, config.internal_expand_px, width, height
        )
        near_expanded = _expanded_bbox(
            main_bbox, config.near_component_expand_px, width, height
        )

        for component in components:
            if component.label == main.label:
                continue
            area_ratio = component.area / image_area if image_area else 0.0
            if component.touches_border:
                if _component_center_inside(component, expanded):
                    keep.add(component.label)
                    if area_ratio > config.border_component_max_ratio:
                        large_border_kept = True
                else:
                    removed_border += 1
                continue
            if _should_keep_detail(
                component,
                expanded=expanded,
                near_expanded=near_expanded,
                config=config,
            ):
                keep.add(component.label)

    filtered = np.zeros_like(raw_mask)
    if keep:
        filtered[np.isin(labels, list(keep))] = 255

    # Только closing — без opening/эрозии, чтобы не рвать тонкие дужки.
    kernel_size = max(1, config.morphology_kernel_size)
    if kernel_size % 2 == 0:
        kernel_size += 1
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (kernel_size, kernel_size))
    closed = cv2.morphologyEx(
        filtered,
        cv2.MORPH_CLOSE,
        kernel,
        iterations=max(0, config.morphology_iterations),
    )

    labels2, components2 = _analyze_components(closed)
    keep2: set[int] = set()
    significant = [c for c in components2 if c.area >= config.min_component_area]
    if significant:
        non_border = [c for c in significant if not c.touches_border]
        central = [
            c
            for c in significant
            if (not c.touches_border) or _center_of_mass_in_central_zone(c)
        ]
        pool = non_border or central or significant
        main2 = max(pool, key=lambda item: item.area)
        keep2.add(main2.label)
        expanded2 = _expanded_bbox(
            (main2.x, main2.y, main2.width, main2.height),
            config.internal_expand_px,
            width,
            height,
        )
        near2 = _expanded_bbox(
            (main2.x, main2.y, main2.width, main2.height),
            config.near_component_expand_px,
            width,
            height,
        )
        for component in components2:
            if component.label == main2.label:
                continue
            if component.touches_border:
                if _component_center_inside(component, expanded2):
                    keep2.add(component.label)
                continue
            if _should_keep_detail(
                component,
                expanded=expanded2,
                near_expanded=near2,
                config=config,
            ):
                keep2.add(component.label)

    clean = np.zeros_like(closed)
    if keep2:
        clean[np.isin(labels2, list(keep2))] = 255

    clean_ratio = float(np.count_nonzero(clean) / image_area) if image_area else 0.0
    _final_labels, final_components = _analyze_components(clean)
    significant_count = sum(
        1 for c in final_components if c.area >= config.min_component_area
    )

    non_border_raw = _non_border_mask(raw_mask)
    non_border_pixels = int(np.count_nonzero(non_border_raw))
    clean_pixels = int(np.count_nonzero(clean))
    removed_object_pixels = int(np.count_nonzero((non_border_raw > 0) & (clean == 0)))
    if non_border_pixels > 0:
        retention = float(np.count_nonzero((non_border_raw > 0) & (clean > 0))) / float(
            non_border_pixels
        )
    else:
        retention = 1.0 if clean_pixels == 0 else 0.0
    loss = 1.0 - retention

    reference = structure_reference if structure_reference is not None else raw_mask
    reference_nb = _non_border_mask(reference)
    ref_labels, ref_components = _analyze_components(reference_nb)
    clean_bbox = _mask_bbox(clean)
    if ref_components and clean_bbox is not None:
        near_ref = _expanded_bbox(
            clean_bbox, config.near_component_expand_px, width, height
        )
        dilate_px = max(3, config.thin_reclaim_dilate_px)
        dilate_k = cv2.getStructuringElement(
            cv2.MORPH_ELLIPSE, (dilate_px * 2 + 1, dilate_px * 2 + 1)
        )
        clean_dilated = cv2.dilate(clean, dilate_k, iterations=2)
        related: list[ComponentInfo] = []
        for component in ref_components:
            if component.area < config.min_elongated_component_area:
                continue
            if not _component_near_or_inside(component, near_ref):
                continue
            comp_pixels = ref_labels == component.label
            overlaps_clean = bool(np.any(clean_dilated[comp_pixels] > 0))
            elongated = _is_elongated(component, config)
            if overlaps_clean or elongated:
                related.append(component)
        if not related:
            related = [max(ref_components, key=lambda item: item.area)]
        x0 = min(c.x for c in related)
        y0 = min(c.y for c in related)
        x1 = max(c.x + c.width - 1 for c in related)
        y1 = max(c.y + c.height - 1 for c in related)
        ref_bbox = (x0, y0, x1 - x0 + 1, y1 - y0 + 1)
    elif ref_components:
        main_ref = max(ref_components, key=lambda item: item.area)
        ref_bbox = (main_ref.x, main_ref.y, main_ref.width, main_ref.height)
    else:
        ref_bbox = None

    if ref_bbox is None or clean_bbox is None:
        extent_retention = 0.0 if ref_bbox is not None else 1.0
    else:
        width_ret = clean_bbox[2] / float(max(1, ref_bbox[2]))
        height_ret = clean_bbox[3] / float(max(1, ref_bbox[3]))
        extent_retention = float(min(1.0, width_ret, height_ret))

    removed_component_count = max(0, initial_component_count - len(final_components))

    return SegmentationResult(
        raw_mask=raw_mask,
        clean_mask=clean,
        raw_foreground_ratio=raw_ratio,
        clean_foreground_ratio=clean_ratio,
        component_count=len(final_components),
        significant_component_count=significant_count,
        border_component_count=border_component_count,
        removed_border_component_count=removed_border,
        removed_component_count=removed_component_count,
        removed_foreground_pixels=removed_object_pixels,
        foreground_retention_ratio=retention,
        foreground_loss_ratio=loss,
        object_extent_retention_ratio=extent_retention,
        large_border_touch_kept=large_border_kept,
        components=final_components,
    )


def segment_image(
    image: Image.Image,
    background: BackgroundStats,
    config: MaskConfig,
) -> SegmentationResult:
    """Построить сырую маску и очистить её."""
    raw = build_raw_mask(image, background, config)
    probe = build_structure_probe_mask(image, background, config)
    return clean_mask(raw, config, structure_reference=probe)