docreader-ocr 0.2.2__tar.gz → 0.2.4__tar.gz

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: docreader-ocr
-Version: 0.2.2
+Version: 0.2.4
 Summary: Document OCR pipeline: classify → detect fields → recognize text
 Project-URL: Homepage, https://github.com/mishanyacorleone/docreader
 Project-URL: Repository, https://github.com/mishanyacorleone/docreader
@@ -18,6 +18,7 @@ Requires-Python: >=3.9
 Requires-Dist: easyocr>=1.7
 Requires-Dist: numpy>=1.24
 Requires-Dist: opencv-python>=4.8
+Requires-Dist: rapidfuzz>=3.14.0
 Requires-Dist: requests>=2.28
 Requires-Dist: torch>=2.0
 Requires-Dist: torchvision>=0.15

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "docreader-ocr"
-version = "0.2.2"
+version = "0.2.4"
 description = "Document OCR pipeline: classify → detect fields → recognize text"
 readme = "README.md"
 license = {text = "MIT"}
@@ -32,6 +32,7 @@ dependencies = [
     "numpy>=1.24",
     "requests>=2.28",
     "tqdm>=4.65",
+    "RapidFuzz>=3.14.0"
 ]
 
 [project.optional-dependencies]

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/src/docreader/config.py

@@ -6,6 +6,13 @@ from dataclasses import dataclass, field
 
 DEFAULT_SKIP_OCR_ZONES = frozenset({"stamp", "gerb"})
 
+DEFAULT_AMBIGUOUS_CLASSES = frozenset({"attestat/diplom"})
+
+DEFAULT_SUBTYPE_KEYWORDS: dict[str, list[str]] = {
+    "attestat": ["аттестат"],
+    "diplom": ["диплом"]
+}
+
 
 @dataclass
 class PipelineConfig:
@@ -35,6 +42,18 @@ class PipelineConfig:
     ocr_download_enabled: bool = False
     skip_ocr_zones: frozenset[str] = DEFAULT_SKIP_OCR_ZONES 
 
+    # Resolver
+    ambiguous_classes: frozenset[str] = field(
+        default_factory=lambda: DEFAULT_AMBIGUOUS_CLASSES
+    )
+    resolver_weights: str = "lvl_detector.pt"
+    resolver_confidence: float = 0.25
+    resolver_subtype_keywords: dict[str, list[str]] = field(
+        default_factory=lambda: dict(DEFAULT_SUBTYPE_KEYWORDS)
+    )
+    resolver_fuzzy_threshold: float = 60.0
+    resolver_fallback: str | None = None
+
     enable_deskew: bool = True  # Выравнивание по линиям Хафа
     return_crops: bool = True  # Сохранять кропы зон в результат
 

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/src/docreader/factory.py

@@ -19,6 +19,7 @@ from docreader.hub import ensure_model
 from docreader.classifier.yolo_classifier import DocClassifier
 from docreader.detector.yolo_obb import ZoneDetector
 from docreader.ocr.easyocr_engine import TextRecognizer
+from docreader.resolver.lvl_resolver import LvlSubtypeResolver
 
 
 def create_classifier(
@@ -127,3 +128,41 @@ def create_ocr(
     }
     defaults.update(kwargs)
     return TextRecognizer(**defaults)
+
+
+def create_resolver(
+    config: PipelineConfig | None = None,
+    ocr_engine: TextRecognizer | None = None,
+    **kwargs
+) -> LvlSubtypeResolver:
+    """
+    Создаёт resolver подтипа документа (attestat/diplom)
+
+    Args:
+        config: конфигурация (если None — используется дефолтная).
+        ocr_engine: готовый OCR-движок (если None - создаётся новый).
+        **kwargs: переопределение параметров TextRecognizer
+            (weights_path, match_threshold, detector_confidence, device).
+
+    Returns:
+        Готовый к работе LvlSubtypeResolver.
+
+    Примеры:
+        resolver = create_resolver()
+        resolver = create_resolver(match_threshold=70.0)
+        resolver = create_resolver(ocr_engine=my_ocr)
+    """
+    cfg = config or PipelineConfig()
+    ocr = ocr_engine or create_ocr(cfg)
+
+    defaults = {
+        "weights_path": cfg.resolver_weights,
+        "ocr_engine": ocr_engine,
+        "subtype_keywords": cfg.resolver_subtype_keywords,
+        "fuzzy_threshold": cfg.resolver_fuzzy_threshold,
+        "confidence_threshold": cfg.resolver_confidence,
+        "fallback": cfg.resolver_fallback,
+        "device": cfg.resolve_device(),
+    }
+    defaults.update(kwargs)
+    return LvlSubtypeResolver(**defaults)

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/src/docreader/hub.py

@@ -17,6 +17,7 @@ from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
+_BASE_LVL_DETECTOR = "https://github.com/mishanyacorleone/docreader/releases/download/v0.2.2"
 _BASE_URL_CLASSIFIER = "https://github.com/mishanyacorleone/docreader/releases/download/v0.2.1"
 _BASE_URL = "https://github.com/mishanyacorleone/docreader/releases/download/v0.1.0"
 
@@ -24,8 +25,8 @@ MODEL_REGISTRY: dict[str, dict] = {
     # === Классификатор документов (YOLO OBB) ===
     "doc_classifier.pt": {
         "url": f"{_BASE_URL_CLASSIFIER}/doc_classifier.pt",
-        "sha256": "d912884d8517cf776e989dc4fced855f34c2ee1d8b17732b778d7e84b7de84fc",
-        "size_mb": 6.03,
+        "sha256": "b97567bfe70a13974219a0cf181643f322868054bac9c6e9ff812eb6ec2607b2",
+        "size_mb": 5.68,
     },
 
     # === Детекторы зон ===
@@ -56,7 +57,14 @@ MODEL_REGISTRY: dict[str, dict] = {
         "sha256": "832ce5a7f3a1086d81beb1c991347e3f545a425646bc87f3f576ae06fecd2420",
         "size_mb": 87.1,
         "extract_to": "easyocr"
-    }
+    },
+
+    # === Resolver ===
+    "lvl_detector.pt": {
+        "url": f"{_BASE_LVL_DETECTOR}/lvl_detector.pt",
+        "sha256": "10bc71dbf8de891bc591154c3c369d8db2daa329249ef4e5b4b15508e8441ba4",
+        "size_mb": 5.63,
+    },
 }
 
 def get_cache_dir() -> Path:

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/src/docreader/pipeline.py

@@ -14,10 +14,9 @@ from docreader.hub import ensure_model
 from docreader.preprocessing import deskew_image, crop_obb_region
 
 from docreader.classifier.base import BaseClassifier
-
 from docreader.detector.base import BaseDetector
-
 from docreader.ocr.base import BaseOcrEngine
+from docreader.resolver.base import BaseSubtypeResolver
 
 logger = logging.getLogger(__name__)
 
@@ -47,6 +46,7 @@ class DocReader:
         classifier: Optional[BaseClassifier] = None,
         detector: Optional[BaseDetector] = None,
         ocr_engine: Optional[BaseOcrEngine] = None,
+        subtype_resolver: Optional[BaseSubtypeResolver] = None
     ):
         self._config = config or PipelineConfig()
         self._device = self._config.resolve_device()
@@ -56,6 +56,7 @@ class DocReader:
         self._classifier = classifier or self._build_classifier()
         self._detector = detector or self._build_detector()
         self._ocr = ocr_engine or self._build_ocr()
+        self._resolver = subtype_resolver or self._build_resolver()
 
     def _build_classifier(self) -> BaseClassifier:
         """Создаёт классификатор из конфига."""
@@ -101,6 +102,28 @@ class DocReader:
             recog_network=self._config.ocr_recog_network,
             download_enabled=self._config.ocr_download_enabled,
         )
+    
+    def _build_resolver(self) -> Optional[BaseSubtypeResolver]:
+        """
+        Создаёт resolver только если есть неоднозначные классы в конфиге.
+        Возвращает None, если resolver не нужен
+        """
+        if not self._config.ambiguous_classes:
+            return None
+        
+        from docreader.resolver.lvl_resolver import LvlSubtypeResolver
+
+        weights_path = ensure_model(self._config.resolver_weights)
+
+        return LvlSubtypeResolver(
+            weights_path=weights_path,
+            ocr_engine=self._ocr,
+            subtype_keywords=self._config.resolver_subtype_keywords,
+            fuzzy_threshold=self._config.resolver_fuzzy_threshold,
+            confidence_threshold=self._config.resolver_confidence,
+            fallback=self._config.resolver_fallback,
+            device=self._device,
+        )
 
     # === Публичный API ===
 
@@ -111,11 +134,11 @@ class DocReader:
     ) -> PageResult:
         """
         Полный пайплайн: находит все документы и распознаёт.
-
+ 
         Args:
             source: путь к файлу или numpy array (BGR).
             return_crops: сохранять ли кропы.
-
+ 
         Returns:
             PageResult со списком найденных документов.
         """
@@ -124,17 +147,14 @@ class DocReader:
             if return_crops is not None
             else self._config.return_crops
         )
-
+ 
         image = load_image(source)
-
-        # 1. Классификация
         classified_docs = self._classifier.classify(image)
-
+ 
         if not classified_docs:
             logger.info("No documents found")
             return PageResult(documents=[])
-
-        # 2. Обработка каждого документа
+ 
         documents: list[DocumentResult] = []
         for doc in classified_docs:
             result = self._process_single_document(
@@ -145,11 +165,11 @@ class DocReader:
                 save_crops=save_crops,
             )
             documents.append(result)
-
+ 
         page_result = PageResult(documents=documents)
         logger.info(f"Complete: {page_result}")
         return page_result
-
+ 
     def process_batch(
         self,
         sources: list[ImageSource],
@@ -159,6 +179,45 @@ class DocReader:
         return [self.process(src, return_crops) for src in sources]
 
     # === Внутренняя логика ===
+    
+    def _resolve_doc_type(
+        self,
+        doc_type: str,
+        doc_image: np.ndarray
+    ) -> tuple[str, dict]:
+        """
+        Уточняет тип документа через resolver, если класс неоднозначен.
+
+        Returns:
+            Кортеж (уточнённый doc_type, метаданные resolve).
+        """
+        if (
+            doc_type not in self._config.ambiguous_classes
+            or self._resolver is None
+        ):
+            return doc_type, {}
+
+        resolve_result = self._resolver.resolve(doc_image)
+        meta = {
+            "resolver_ocr_text": resolve_result.ocr_text,
+            "resolver_ocr_confidence": resolve_result.confidence,
+            "resolver_fuzzy_score": resolve_result.fuzzy_score
+        }
+
+        if resolve_result.resolve:
+            logger.info(
+                f"Resolved '{doc_type}' -> '{resolve_result.subtype}' "
+                f"(text='{resolve_result.ocr_text}', "
+                f"fuzzy={resolve_result.fuzzy_score:.1f})"
+            )
+            return resolve_result.subtype, meta
+        
+        logger.warning(
+            f"Could not resolve subtype for '{doc_type}': "
+            f"text='{resolve_result.ocr_text}', "
+            f"score={resolve_result.fuzzy_score:.1f}"
+        )
+        return doc_type, meta
 
     def _process_single_document(
         self,
@@ -169,42 +228,44 @@ class DocReader:
         save_crops: bool,
     ) -> DocumentResult:
         """Обрабатывает один документ."""
-
-        if doc_type not in self._detector.supported_doc_types:
-            logger.warning(f"No detector for '{doc_type}'")
+        if self._config.enable_deskew:
+            doc_image = deskew_image(doc_image)
+ 
+        resolved_type, resolve_meta = self._resolve_doc_type(doc_type, doc_image)
+ 
+        if resolved_type not in self._detector.supported_doc_types:
+            logger.warning(f"No detector for '{resolved_type}'")
             return DocumentResult(
-                doc_type=doc_type,
+                doc_type=resolved_type,
                 doc_confidence=doc_confidence,
                 zones=[],
                 doc_bbox=doc_bbox.tolist(),
                 doc_crop=doc_image if save_crops else None,
+                resolve_meta=resolve_meta,
             )
-
-        if self._config.enable_deskew:
-            doc_image = deskew_image(doc_image)
-
-        detections = self._detector.detect(doc_image, doc_type)
-        logger.info(f"'{doc_type}': {len(detections)} zones")
-
+ 
+        detections = self._detector.detect(doc_image, resolved_type)
+        logger.info(f"'{resolved_type}': {len(detections)} zones")
+ 
         zones: list[ZoneResult] = []
         for det in detections:
             zone = self._process_zone(doc_image, det, save_crops)
             if zone is not None:
                 zones.append(zone)
-
+ 
         return DocumentResult(
-            doc_type=doc_type,
+            doc_type=resolved_type,
             doc_confidence=doc_confidence,
             zones=zones,
             doc_bbox=doc_bbox.tolist(),
             doc_crop=doc_image if save_crops else None,
+            resolve_meta=resolve_meta,
         )
 
     def _process_zone(self, image, detection, save_crops):
         """Обрабатывает одну зону."""
-
         zone_name = detection.zone_name
-
+ 
         if zone_name in self._config.skip_ocr_zones:
             return ZoneResult(
                 name=zone_name,
@@ -212,14 +273,14 @@ class DocReader:
                 confidence=detection.confidence,
                 bbox=detection.obb_points.tolist(),
             )
-
+ 
         crop = crop_obb_region(image, detection.obb_points)
         if crop is None or crop.size == 0:
             logger.warning(f"Empty crop for '{zone_name}'")
             return None
-
+ 
         ocr_result = self._ocr.recognize(crop)
-
+ 
         return ZoneResult(
             name=zone_name,
             text=ocr_result.text,
@@ -239,6 +300,7 @@ class DocReader:
         self._classifier = None
         self._detector = None
         self._ocr = None
+        self._resolver = None
         try:
             import gc
             gc.collect()

docreader_ocr-0.2.4/src/docreader/resolver/__init__.py

@@ -0,0 +1,4 @@
+from docreader.resolver.base import BaseSubtypeResolver, ResolveResult
+from docreader.resolver.lvl_resolver import LvlSubtypeResolver
+
+__all__ = ["BaseSubtypeResolver", "ResolveResult", "LvlSubtypeResolver"]

docreader_ocr-0.2.4/src/docreader/resolver/base.py

@@ -0,0 +1,40 @@
+"""
+Абстрактный интерфейс resolver'a подтипа документа.
+"""
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional
+
+import numpy as np
+
+
+@dataclass
+class ResolveResult:
+    """
+    Результат определения подтипа документа
+    """
+    subtype: Optional[str]
+    ocr_text: str
+    confidence: float
+    fuzzy_score: float
+
+    @property
+    def resolve(self) -> None:
+        return self.subtype is not None
+
+
+class BaseSubtypeResolver(ABC):
+    """
+    Интерфейс для определения подтипа документа.
+
+    Используется когда классификатор не может различать 2 похожих
+    класса (attestat/diplom) и требуется дополнительный шаг
+    """
+
+    @abstractmethod
+    def resolve(self, image: np.ndarray) -> ResolveResult:
+        """
+        Определяет подтип документа по его crop'y
+        """
+        ...

docreader_ocr-0.2.4/src/docreader/resolver/lvl_resolver.py

@@ -0,0 +1,195 @@
+"""
+Resolver подтипа документа через детекцию поля lvl + OCR + fuzzy matching.
+"""
+
+import logging
+from typing import Optional
+
+import numpy as np
+from rapidfuzz import process, fuzz
+from ultralytics import YOLO
+
+from docreader.ocr.base import BaseOcrEngine
+from docreader.preprocessing.geometry import crop_obb_region
+from docreader.resolver.base import BaseSubtypeResolver, ResolveResult
+
+logger = logging.getLogger(__name__)
+
+
+class LvlSubtypeResolver(BaseSubtypeResolver):
+    """
+    Определяет подтип документа (attestat/diplom) через:
+      1. YOLO OBB — детектирует поле lvl на crop'е документа
+      2. OCR — распознаёт текст поля
+      3. Fuzzy matching — сопоставляет текст с ключевыми словами подтипов
+
+    Примеры:
+        resolver = LvlSubtypeResolver(
+            weights_path="/path/to/lvl_detector.pt",
+            ocr_engine=ocr,
+            subtype_keywords={
+                "attestat": ["аттестат", "attestat"],
+                "diplom": ["диплом", "diplom"],
+            },
+        )
+        result = resolver.resolve(doc_crop)
+        if result.resolved:
+            print(result.subtype)  # "attestat" или "diplom"
+
+    Args:
+        weights_path: путь к YOLO-модели для детекции поля lvl.
+        ocr_engine: движок OCR (BaseOcrEngine).
+        subtype_keywords: словарь {подтип: [ключевые слова]}.
+        fuzzy_threshold: минимальный score для признания совпадения (0–100).
+        confidence_threshold: минимальная уверенность детектора.
+        fallback: подтип по умолчанию если resolve не удался (None = unresolved).
+        device: устройство ("cpu", "cuda").
+    """
+
+    def __init__(
+        self,
+        weights_path: str,
+        ocr_engine: BaseOcrEngine,
+        subtype_keywords: dict[str, list[str]],
+        fuzzy_threshold: float = 60.0,
+        confidence_threshold: float = 0.25,
+        fallback: Optional[str] = None,
+        device: str = "cpu",
+    ):
+        self._model = YOLO(weights_path)
+        self._ocr = ocr_engine
+        self._fuzzy_threshold = fuzzy_threshold
+        self._confidence_threshold = confidence_threshold
+        self._fallback = fallback
+        self._device = device
+
+        # Плоский список ключевых слов и маппинг слово → подтип
+        self._keywords: list[str] = []
+        self._keyword_to_subtype: dict[str, str] = {}
+        for subtype, words in subtype_keywords.items():
+            for word in words:
+                normalized = word.lower()
+                self._keywords.append(normalized)
+                self._keyword_to_subtype[normalized] = subtype
+
+        logger.info(
+            f"LvlSubtypeResolver initialized: "
+            f"subtypes={list(subtype_keywords.keys())}, "
+            f"threshold={fuzzy_threshold}, fallback={fallback}"
+        )
+
+    def resolve(self, image: np.ndarray) -> ResolveResult:
+        """
+        Определяет подтип документа по crop'у.
+
+        Args:
+            image: BGR изображение документа.
+
+        Returns:
+            ResolveResult с подтипом и диагностической информацией.
+        """
+        lvl_crop = self._detect_lvl_field(image)
+
+        if lvl_crop is None:
+            logger.warning("lvl field not detected, using fallback")
+            return ResolveResult(
+                subtype=self._fallback,
+                ocr_text="",
+                confidence=0.0,
+                fuzzy_score=0.0,
+            )
+
+        ocr_result = self._ocr.recognize(lvl_crop)
+        logger.debug(
+            f"lvl OCR: text='{ocr_result.text}', conf={ocr_result.confidence:.3f}"
+        )
+
+        if not ocr_result.text.strip():
+            logger.warning("lvl OCR returned empty text, using fallback")
+            return ResolveResult(
+                subtype=self._fallback,
+                ocr_text=ocr_result.text,
+                confidence=ocr_result.confidence,
+                fuzzy_score=0.0,
+            )
+
+        subtype, fuzzy_score = self._match_subtype(ocr_result.text)
+
+        if subtype is None:
+            logger.warning(
+                f"Fuzzy match below threshold: "
+                f"text='{ocr_result.text}', score={fuzzy_score:.1f}, "
+                f"threshold={self._fuzzy_threshold}, fallback={self._fallback}"
+            )
+
+        return ResolveResult(
+            subtype=subtype,
+            ocr_text=ocr_result.text,
+            confidence=ocr_result.confidence,
+            fuzzy_score=fuzzy_score,
+        )
+
+    def _detect_lvl_field(self, image: np.ndarray) -> Optional[np.ndarray]:
+        """
+        Детектирует поле lvl и возвращает его crop.
+        """
+        
+        results = self._model(image, device=self._device, verbose=False)
+
+        if results[0].obb is None:
+            return None
+
+        best_conf = -1.0
+        best_crop = None
+
+        for det in results[0].obb:
+            confidence = float(det.conf.cpu())
+            if confidence < self._confidence_threshold:
+                continue
+
+            zone_name = self._model.names[int(det.cls.cpu())]
+            if zone_name != "lvl":
+                continue
+
+            if confidence <= best_conf:
+                continue
+
+            obb_points = det.xyxyxyxy.cpu().numpy().flatten()
+            crop = crop_obb_region(image, obb_points)
+
+            if crop is not None and crop.size > 0:
+                best_conf = confidence
+                best_crop = crop
+
+        return best_crop
+
+    def _match_subtype(self, text: str) -> tuple[Optional[str], float]:
+        """
+        Сопоставляет OCR-текст с ключевыми словами через fuzzy matching.
+
+        Returns:
+            Кортеж (подтип или None, fuzzy score).
+        """
+        normalized = text.lower().strip()
+
+        match = process.extractOne(
+            normalized,
+            self._keywords,
+            scorer=fuzz.WRatio,
+        )
+
+        if match is None:
+            return None, 0.0
+
+        best_keyword, score, _ = match
+
+        if score < self._fuzzy_threshold:
+            return None, float(score)
+
+        subtype = self._keyword_to_subtype[best_keyword]
+        logger.debug(
+            f"Fuzzy matched: '{normalized}' -> '{best_keyword}' "
+            f"(subtype={subtype}, score={score:.1f})"
+        )
+        return subtype, float(score)
+    

{docreader_ocr-0.2.2 → docreader_ocr-0.2.4}/src/docreader/schemas.py

@@ -33,6 +33,7 @@ class DocumentResult:
     zones: list[ZoneResult] = field(default_factory=list)
     doc_bbox: Optional[list[float]] = None # координаты документа в исходном изображении
     doc_crop: Optional[np.ndarray] = None # кроп документа
+    resolve_meta: dict = field(default_factory=dict) # диагностика resolver'a
 
     @property
     def fields(self) -> dict[str, str]:
@@ -42,7 +43,7 @@ class DocumentResult:
         return {zone.name: zone.text for zone in self.zones}
     
     def to_dict(self) -> dict:
-        return {
+        result = {
             "document": {
                 "doc_type": self.doc_type,
                 "doc_confidence": round(self.doc_confidence, 4),
@@ -50,6 +51,9 @@ class DocumentResult:
                 "fields": self.fields
             }
         }
+        if self.resolve_meta:
+            result["document"]["resolve_meta"] = self.resolve_meta
+        return result
     
     def __repr__(self) -> str:
         return (