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

src/photo_processor/background.py
122 строки4 KB
otcheskiy
feat: add simple background segmentation and mask quality
01 авг 2026, 12:34
01 авг 2026, 12:3403c5abe
"""Оценка светлого фона по приграничным областям."""

from __future__ import annotations

import logging
from dataclasses import dataclass

import numpy as np
from PIL import Image

from photo_processor.config import BackgroundConfig

logger = logging.getLogger(__name__)


class BackgroundError(ValueError):
    """Не удалось устойчиво оценить светлый фон."""


@dataclass(frozen=True)
class BackgroundStats:
    median_rgb: tuple[float, float, float]
    mean_rgb: tuple[float, float, float]
    brightness: float
    brightness_std: float
    usable_ratio: float
    sample_count: int


def _luminance(rgb: np.ndarray) -> np.ndarray:
    """Яркость в приближении Rec. 601, rgb shape (..., 3)."""
    return (
        0.299 * rgb[..., 0]
        + 0.587 * rgb[..., 1]
        + 0.114 * rgb[..., 2]
    )


def _collect_border_samples(rgb: np.ndarray, border_ratio: float) -> np.ndarray:
    height, width, _ = rgb.shape
    border_x = max(1, int(round(width * border_ratio)))
    border_y = max(1, int(round(height * border_ratio)))
    strips = [
        rgb[:border_y, :, :],
        rgb[-border_y:, :, :],
        rgb[:, :border_x, :],
        rgb[:, -border_x:, :],
    ]
    return np.concatenate([strip.reshape(-1, 3) for strip in strips], axis=0)


def estimate_background(
    image: Image.Image,
    config: BackgroundConfig,
) -> BackgroundStats:
    """Оценить цвет и яркость светлого фона по границам кадра."""
    if image.mode != "RGB":
        raise BackgroundError(
            f"Ожидалось RGB-изображение для оценки фона, получено: {image.mode}."
        )

    rgb = np.asarray(image, dtype=np.float32)
    if rgb.ndim != 3 or rgb.shape[2] != 3:
        raise BackgroundError("Некорректная форма массива изображения для оценки фона.")

    samples = _collect_border_samples(rgb, config.sample_border_ratio)
    if samples.size == 0:
        raise BackgroundError("Не удалось набрать пиксели границ для оценки фона.")

    luminances = _luminance(samples)
    dark_cutoff = float(np.percentile(luminances, config.reject_dark_percentile))
    usable_mask = luminances >= max(dark_cutoff, 0.0)
    usable = samples[usable_mask]
    usable_ratio = float(usable_mask.mean()) if usable_mask.size else 0.0

    if usable.shape[0] == 0:
        raise BackgroundError(
            "После исключения тёмных пикселей не осталось пригодных фоновых выборок."
        )

    if usable_ratio < config.min_usable_ratio:
        raise BackgroundError(
            "Слишком мала доля пригодных фоновых пикселей на границах: "
            f"{usable_ratio:.4f} < {config.min_usable_ratio}."
        )

    usable_lum = _luminance(usable)
    brightness = float(np.median(usable_lum))
    brightness_std = float(np.std(usable_lum))
    median_rgb = tuple(float(v) for v in np.median(usable, axis=0))
    mean_rgb = tuple(float(v) for v in np.mean(usable, axis=0))

    if brightness < config.min_brightness:
        raise BackgroundError(
            "Оценённый фон недостаточно светлый: "
            f"brightness={brightness:.1f} < {config.min_brightness}."
        )
    if brightness_std > config.max_brightness_std:
        raise BackgroundError(
            "Слишком большой разброс яркости фона: "
            f"std={brightness_std:.1f} > {config.max_brightness_std}."
        )

    stats = BackgroundStats(
        median_rgb=median_rgb,
        mean_rgb=mean_rgb,
        brightness=brightness,
        brightness_std=brightness_std,
        usable_ratio=usable_ratio,
        sample_count=int(usable.shape[0]),
    )
    logger.debug(
        "Фон: rgb_median=(%.1f, %.1f, %.1f) brightness=%.1f std=%.1f usable=%.3f n=%d",
        stats.median_rgb[0],
        stats.median_rgb[1],
        stats.median_rgb[2],
        stats.brightness,
        stats.brightness_std,
        stats.usable_ratio,
        stats.sample_count,
    )
    return stats