app/core/document.py - draft-ui - ksrk/layout_validation

app/core/document.py
276 строк10 KB
ksrk
[ui] minor improvement: introduced interactive layer view in the streamlit interface
07 авг 2026, 23:28
07 авг 2026, 23:28276aea0
from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from enum import Enum
from typing import Dict, FrozenSet, Iterator, List, Optional, Set

import numpy as np
from pymupdf import Page, Document

from core.utils.pdf import convert_to_cmyk_raster

logger = logging.getLogger(__name__)


class LayerName(str, Enum):
    CUT         = "CUT"
    WHITEWASH   = "WHITEWASH"
    VARNISH     = "VARNISH"
    IMAGES_CMYK = "IMAGE_CMYK"
    PANTONE     = "PANTONE"
    MATERIAL    = "MATERIAL"
    NO_PRINT    = "NO_PRINT"

    @classmethod
    def known_prefixes(cls) -> FrozenSet[str]:
        return frozenset(m.value.upper() for m in cls)


@dataclass
class LayerInfo:
    xref: int
    raw_name: str
    layer_type: Optional[LayerName]
    pages: Set[int]
    visible_by_default: bool

    def applies_to_page(self, page_num: int) -> bool:
        if not self.pages:
            return True
        return page_num in self.pages


def _parse_layer_name(raw_name: str) -> tuple[Optional[LayerName], Set[int]]:
    """
    Разбирает имя слоя на тип и набор страниц.
    Примеры: "CUT_12" -> (LayerName.CUT, {1, 2}), "IMAGE_CMYK" -> (LayerName.IMAGES_CMYK, set())
    """
    upper = raw_name.upper().strip()

    match = re.match(r'^(.+?)_(\d+)$', upper)
    if match:
        prefix = match.group(1).strip()
        digits = match.group(2)
        pages = {int(d) for d in digits}
    else:
        prefix = upper
        pages = set()

    layer_type = None
    for known in LayerName.known_prefixes():
        if prefix == known:
            layer_type = LayerName(known)
            break

    return layer_type, pages


class DocumentWrapper:
    def __init__(self, doc: Document, filename: str) -> None:
        self._name = filename
        self._doc = doc
        self._layers: Dict[int, LayerInfo] = {}
        self._by_type: Dict[LayerName, List[LayerInfo]] = {}
        self._scan_layers()

    @property
    def doc(self) -> Document:
        return self._doc

    @property
    def filename(self) -> str:
        return self._name

    @property
    def page_count(self) -> int:
        return len(self._doc)

    @property
    def layers(self) -> List[LayerInfo]:
        return list(self._layers.values())

    def has_layer_type(self, layer_type: LayerName) -> bool:
        return layer_type in self._by_type and bool(self._by_type[layer_type])

    def get_layers_of_type(self, layer_type: LayerName) -> List[LayerInfo]:
        return self._by_type.get(layer_type, [])

    def layers_for_page(self, page_num: int) -> List[LayerInfo]:
        return [l for l in self._layers.values() if l.applies_to_page(page_num)]

    def check_required_layers(self, required: List[LayerName]) -> Dict[LayerName, bool]:
        result = {}
        for layer_type in required:
            present = self.has_layer_type(layer_type)
            result[layer_type] = present
            if not present:
                logger.warning(f"Обязательный слой [{layer_type.value}] отсутствует.")
            else:
                layers = self.get_layers_of_type(layer_type)
                logger.info(f"Слой [{layer_type.value}] найден: {[l.raw_name for l in layers]}")
        return result

    def set_only(self, layer_types: List[LayerName], page_num: int) -> None:
        """
        Переключает видимость слоёв напрямую по конфигам,
        корректно обрабатывая дубликаты слоёв с одинаковыми именами.
        """
        active_types = set(layer_types)

        for ui_config in self._doc.layer_ui_configs():
            number = ui_config['number']
            raw_name = ui_config['text']

            layer_type, pages = _parse_layer_name(raw_name)
            applies = not pages or page_num in pages

            visible = (layer_type in active_types) and applies
            self._doc.set_layer_ui_config(number, action=0 if visible else 2)

    def set_visibility_for_page(self, page_num: int) -> None:
        for ui_config in self._doc.layer_ui_configs():
            number = ui_config['number']
            raw_name = ui_config['text']

            _, pages = _parse_layer_name(raw_name)
            visible = not pages or page_num in pages

            self._doc.set_layer_ui_config(number, action=0 if visible else 2)

    def restore_defaults(self) -> None:
        for ui_config in self._doc.layer_ui_configs():
            number = ui_config['number']
            raw_name = ui_config['text']

            layer_type, _ = _parse_layer_name(raw_name)
            visible = True
            if layer_type and layer_type in self._by_type:
                visible = self._by_type[layer_type][0].visible_by_default

            self._doc.set_layer_ui_config(number, action=0 if visible else 2)

    def show_layers(self) -> Dict[int, Dict[str, np.ndarray]]:
        """
        Изолированно рендерит каждый слой страницы в CMYK-растр для Streamlit UI.
        """
        layers_vis = {}
        for page_num in range(self.page_count):
            page_num_1 = page_num + 1
            layers_vis[page_num_1] = {}

            page_layers = self.layers_for_page(page_num_1)

            for target_info in page_layers:
                for ui_config in self._doc.layer_ui_configs():
                    number = ui_config['number']
                    raw_name = ui_config['text']

                    is_target = (raw_name.upper() == target_info.raw_name.upper())
                    self._doc.set_layer_ui_config(number, action=0 if is_target else 2)

                page = self._doc[page_num]
                raster = convert_to_cmyk_raster(page)
                if np.all(raster == 0) or np.all(raster == 255):
                    continue

                layers_vis[page_num_1][target_info.raw_name] = raster

            self.set_visibility_for_page(page_num_1)

        self.restore_defaults()
        return layers_vis

    def apply_layer_visibility(self, visibility_map: Dict[int, bool]) -> None:
        """
        Применяет словарь видимости {number_ocg: True/False} к документу.
        0 = показать (ON), 2 = скрыть (OFF)
        """
        for ui_config in self._doc.layer_ui_configs():
            number = ui_config['number']
            if number in visibility_map:
                visible = visibility_map[number]
                self._doc.set_layer_ui_config(number, action=0 if visible else 2)

    def iter_pages(
            self,
            layer_filter: Optional[List[LayerName]] = None,
    ) -> Iterator[tuple[int, Page, List[LayerInfo]]]:
        try:
            for page_num_0 in range(self.page_count):
                page_num_1 = page_num_0 + 1

                if layer_filter is not None:
                    self.set_only(layer_filter, page_num_1)
                else:
                    self.set_visibility_for_page(page_num_1)

                active = self.layers_for_page(page_num_1)
                logger.debug(f"Страница {page_num_1}: активные слои = {[l.raw_name for l in active]}")
                yield page_num_0, self._doc[page_num_0], active
        finally:
            self.restore_defaults()

    def iter_pages_by_layer_type(
            self,
            layer_type: LayerName,
    ) -> Iterator[tuple[int, Page]]:
        layers = self.get_layers_of_type(layer_type)
        if not layers:
            logger.warning(f"Слой [{layer_type.value}] не найден.")
            return

        all_pages: Set[int] = set()
        for info in layers:
            if not info.pages:
                all_pages.update(range(1, self.page_count + 1))
            else:
                all_pages.update(info.pages)

        try:
            for page_num_1 in sorted(all_pages):
                page_num_0 = page_num_1 - 1
                if page_num_0 >= self.page_count:
                    continue
                self.set_only([layer_type], page_num_1)
                yield page_num_0, self._doc[page_num_0]
        finally:
            self.restore_defaults()

    def _scan_layers(self) -> None:
        ocgs = self._doc.get_ocgs()
        if not ocgs:
            logger.info("Документ не содержит OCG-слоёв.")
            return

        for xref, info in ocgs.items():
            raw_name: str = info.get("name", "")
            layer_type, pages = _parse_layer_name(raw_name)
            if layer_type is None:
                logger.debug(f"Пропуск нераспознанного или устаревшего OCG-слоя: '{raw_name}' (xref={xref})")
                continue

            layer_info = LayerInfo(
                xref=xref,
                raw_name=raw_name,
                layer_type=layer_type,
                pages=pages,
                visible_by_default=info.get("on", True),
            )
            self._layers[xref] = layer_info
            self._by_type.setdefault(layer_type, []).append(layer_info)

        self._log_summary()

    def _log_summary(self) -> None:
        logger.info(f"Найдено OCG-слоёв: {len(self._layers)}")
        for info in self._layers.values():
            type_tag = f"[{info.layer_type.value}]" if info.layer_type else "[неизвестный]"
            pages_str = ", ".join(str(p) for p in sorted(info.pages)) if info.pages else "все"
            logger.info(
                f"  {type_tag} '{info.raw_name}' "
                f"xref={info.xref} "
                f"страницы={pages_str} "
                f"visible_default={info.visible_by_default}"
            )