docreader-ocr 0.1.2__tar.gz → 0.2.0__tar.gz

{docreader_ocr-0.1.2 → docreader_ocr-0.2.0}/.gitignore

@@ -63,4 +63,12 @@ htmlcov/
 *.log
 .env
 .env.local
-output/
+output/
+
+# ══════════════════════════════════════════
+# Tests
+# ══════════════════════════════════════════ 
+
+tests/*.jpg
+tests/*.png
+tests/*.jpeg

{docreader_ocr-0.1.2 → docreader_ocr-0.2.0}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: docreader-ocr
-Version: 0.1.2
+Version: 0.2.0
 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
 Project-URL: Issues, https://github.com/mishanyacorleone/docreader/issues
-Author-email: Mikhail Kardash <mishutqac@mail.com>, Ruslan Abzelilov <ruslanr26@mail.ru>, Ekaterina Karmanova <monitor81@mail.ru>
+Author-email: Mikhail Kardash <mishutqac@mail.ru>, Ruslan Abzelilov <ruslanr26@mail.ru>, Ekaterina Karmanova <monitor81@mail.ru>
 License: MIT
 License-File: LICENSE
 Keywords: document,ocr,recognition,yolo

{docreader_ocr-0.1.2 → docreader_ocr-0.2.0}/pyproject.toml

@@ -4,13 +4,13 @@ build-backend = "hatchling.build"
 
 [project]
 name = "docreader-ocr"
-version = "0.1.2"
+version = "0.2.0"
 description = "Document OCR pipeline: classify → detect fields → recognize text"
 readme = "README.md"
 license = {text = "MIT"}
 requires-python = ">=3.9"
 authors = [
-    {name = "Mikhail Kardash", email = "mishutqac@mail.com"},
+    {name = "Mikhail Kardash", email = "mishutqac@mail.ru"},
     {name = "Ruslan Abzelilov", email = "ruslanr26@mail.ru"},
     {name = "Ekaterina Karmanova", email = "monitor81@mail.ru"}
 ]

docreader_ocr-0.2.0/src/docreader/__init__.py

@@ -0,0 +1,65 @@
+"""
+docreader — библиотека для распознавания текста с документов.
+
+Быстрый старт:
+    from docreader import DocReader
+    reader = DocReader()
+    result = reader.process("photo.jpg")
+    for doc in result.documents:
+        print(doc.doc_type, doc.fields)
+
+Standalone-компоненты (с автозагрузкой весов):
+    from docreader import create_classifier, create_detector, create_ocr
+
+    clf = create_classifier()
+    det = create_detector()
+    ocr = create_ocr()
+
+Standalone-компоненты (со своими весами):
+    from docreader import DocClassifier, ZoneDetector, TextRecognizer
+
+    clf = DocClassifier(weights_path="/my/model.pt")
+"""
+
+from docreader.pipeline import DocReader
+from docreader.schemas import DocumentResult, ZoneResult, PageResult
+from docreader.config import PipelineConfig
+
+# Классы компонентов (для кастомных весов)
+from docreader.classifier import DocClassifier, BaseClassifier, ClassifiedDocument
+from docreader.detector import ZoneDetector, BaseDetector, Detection
+from docreader.ocr import TextRecognizer, BaseOcrEngine, OcrResult
+
+# Фабрики (со стандартными весами)
+from docreader.factory import create_classifier, create_detector, create_ocr
+
+__all__ = [
+    # Пайплайн
+    "DocReader",
+    "PipelineConfig",
+
+    # Результаты
+    "PageResult",
+    "DocumentResult",
+    "ZoneResult",
+
+    # Фабрики (стандартные веса)
+    "create_classifier",
+    "create_detector",
+    "create_ocr",
+
+    # Классы компонентов (кастомные веса)
+    "DocClassifier",
+    "ZoneDetector",
+    "TextRecognizer",
+
+    # Базовые классы (для наследования)
+    "BaseClassifier",
+    "ClassifiedDocument",
+    "BaseDetector",
+    "Detection",
+    "BaseOcrEngine",
+    "OcrResult",
+]
+
+__version__ = "0.2.0"

docreader_ocr-0.2.0/src/docreader/classifier/__init__.py

@@ -0,0 +1,4 @@
+from docreader.classifier.base import BaseClassifier, ClassifiedDocument
+from docreader.classifier.yolo_classifier import DocClassifier
+
+__all__ = ["BaseClassifier", "ClassifiedDocument", "DocClassifier"]

{docreader_ocr-0.1.2 → docreader_ocr-0.2.0}/src/docreader/classifier/base.py

@@ -1,9 +1,23 @@
 """Абстрактный интерфейс классификатора документов."""
 
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Union
+
 import numpy as np
 
 
+@dataclass
+class ClassifiedDocument:
+    """
+    Один найденный документ на изображении
+    """
+    doc_type: str
+    confidence: float
+    crop: np.ndarray # вырезанный и выпрямленный документ (BGR)
+    obb_points: np.ndarray # shape (4, 2) - координаты в исходном состоянии
+
+
 class BaseClassifier(ABC):
     """
     Интерфейс для классификатора типа документа.
@@ -13,12 +27,12 @@ class BaseClassifier(ABC):
     """
 
     @abstractmethod
-    def predict(self, image: np.ndarray) -> tuple[str, float]:
+    def classify(self, source: Union[str, np.ndarray]) -> list[str, float]:
         """
         Классифицирует изображение документа.
 
         Args:
-            image: BGR изображение (numpy array).
+            image: путь к файлу или BGR изображение (numpy array).
 
         Returns:
             Кортеж (метка_класса, уверенность).

docreader_ocr-0.2.0/src/docreader/classifier/yolo_classifier.py

@@ -0,0 +1,108 @@
+"""Классификатор документов на основе YOLO OBB."""
+
+import logging
+from typing import Optional, Union
+
+import numpy as np
+from ultralytics import YOLO
+
+from docreader.classifier.base import BaseClassifier, ClassifiedDocument
+from docreader.preprocessing.geometry import crop_obb_region
+from docreader.utils import load_image
+
+logger = logging.getLogger(__name__)
+
+
+class DocClassifier(BaseClassifier):
+    """
+    Классификатор документов через YOLO OBB.
+
+    Находит один или несколько документов на фотографии,
+    определяет тип каждого и вырезает кроп.
+
+    Примеры:
+        # Стандартные веса (путь передаётся из пайплайна или вручную)
+        clf = DocClassifier(weights_path="/path/to/doc_classifier.pt")
+        docs = clf.classify(image)
+
+        # Без параметров — используется через DocReader,
+        # который сам подставит путь из конфига
+        reader = DocReader()
+
+    Args:
+        weights_path: путь к файлу весов YOLO.
+        device: устройство ("cpu", "cuda").
+        confidence_threshold: минимальная уверенность детекции.
+    """
+
+    def __init__(
+        self,
+        weights_path: str,
+        device: str = "cpu",
+        confidence_threshold: float = 0.3,
+    ):
+        self._confidence_threshold = confidence_threshold
+        self._device = device
+        self._model = YOLO(weights_path)
+
+        logger.info(
+            f"DocClassifier initialized: device={self._device}, "
+            f"weights={weights_path}, "
+            f"classes={list(self._model.names.values())}"
+        )
+
+    def classify(self, source: Union[str, np.ndarray]) -> list[ClassifiedDocument]:
+        """
+        Находит документы на изображении.
+
+        Args:
+            image: BGR изображение.
+
+        Returns:
+            Список ClassifiedDocument. Пустой, если документы не найдены.
+        """
+        image = load_image(source)
+        
+        results = self._model(image, device=self._device, verbose=False)
+
+        documents = []
+
+        if results[0].obb is None:
+            logger.info("No documents detected")
+            return documents
+
+        for det in results[0].obb:
+            confidence = float(det.conf.cpu())
+            if confidence < self._confidence_threshold:
+                continue
+
+            class_id = int(det.cls.cpu())
+            doc_type = self._model.names[class_id]
+            obb_points = det.xyxyxyxy.cpu().numpy().reshape(4, 2)
+
+            crop = crop_obb_region(image, obb_points)
+            if crop is None or crop.size == 0:
+                logger.warning(
+                    f"Failed to crop document: type={doc_type}, "
+                    f"conf={confidence:.3f}"
+                )
+                continue
+
+            documents.append(ClassifiedDocument(
+                doc_type=doc_type,
+                confidence=confidence,
+                crop=crop,
+                obb_points=obb_points,
+            ))
+
+        logger.info(
+            f"Found {len(documents)} document(s): "
+            f"{[d.doc_type for d in documents]}"
+        )
+        return documents
+
+    @property
+    def class_names(self) -> list[str]:
+        """Список поддерживаемых типов документов."""
+        return list(self._model.names.values())
+    

{docreader_ocr-0.1.2 → docreader_ocr-0.2.0}/src/docreader/config.py

@@ -14,6 +14,9 @@ class PipelineConfig:
     """
     device: str = "auto"
 
+    classifier_weights: str = "doc_classifier.pt"
+    classifier_confidence: float = 0.3
+
     # Типы документов и пути к YOLO-моделям (относительно models_dir)
     detector_weights: dict[str, str] = field(default_factory=lambda: {
         "attestat": "attestat.pt",
@@ -21,26 +24,26 @@ class PipelineConfig:
         "passport": "passport.pt",
         "snils": "snils.pt",
     })
-
-    # Путь к весам классификатора (относительно models_dir)
-    classification_weights: str = "best_doc_classifier.pth"
-
-    class_labels: list[str] = field(default_factory=lambda: [
-        "attestat", "diplom", "passport", "snils", 'other'
-    ])
+    detector_confidence: float = 0.25
 
     # EasyOCR
-    skip_ocr_zones: frozenset[str] = DEFAULT_SKIP_OCR_ZONES 
     ocr_lang: list[str] = field(default_factory=lambda: ["ru"])
+    ocr_model_archive: str = "easyocr_custom.tar.gz"
+    ocr_model_subdir: str = "model"
+    ocr_network_subdir: str = "user_network"
     ocr_recog_network: str = "custom_example"
     ocr_download_enabled: bool = False
+    skip_ocr_zones: frozenset[str] = DEFAULT_SKIP_OCR_ZONES 
 
     enable_deskew: bool = True  # Выравнивание по линиям Хафа
-
     return_crops: bool = True  # Сохранять кропы зон в результат
 
     def resolve_device(self) -> str:
         if self.device != "auto":
             return self.device
-        import torch
-        return "cuda" if torch.cuda.is_available() else "cpu"
+        try:
+            import torch
+            return "cuda" if torch.cuda.is_available() else "cpu"
+        except ImportError:
+            return "cpu"
+        

docreader_ocr-0.2.0/src/docreader/detector/__init__.py

@@ -0,0 +1,4 @@
+from docreader.detector.base import BaseDetector, Detection
+from docreader.detector.yolo_obb import ZoneDetector
+
+__all__ = ["BaseDetector", "Detection", "ZoneDetector"]

docreader_ocr-0.2.0/src/docreader/detector/yolo_obb.py

@@ -0,0 +1,97 @@
+"""Детектор зон документа на основе YOLO OBB."""
+
+import logging
+
+import numpy as np
+from ultralytics import YOLO
+
+from docreader.detector.base import BaseDetector, Detection
+
+logger = logging.getLogger(__name__)
+
+
+class ZoneDetector(BaseDetector):
+    """
+    Детектор полей документа через YOLO OBB с ленивой загрузкой.
+
+    Примеры:
+        det = ZoneDetector(weights_map={
+            "passport": "/path/to/passport.pt",
+            "diplom": "/path/to/diplom.pt",
+        })
+        zones = det.detect(image, doc_type="passport")
+
+    Args:
+        weights_map: словарь {doc_type: полный_путь_к_весам}.
+        device: устройство ("cpu", "cuda").
+        confidence_threshold: минимальная уверенность.
+    """
+
+    def __init__(
+        self,
+        weights_map: dict[str, str],
+        device: str = "cpu",
+        confidence_threshold: float = 0.25,
+    ):
+        self._weights_map = weights_map
+        self._device = device
+        self._confidence_threshold = confidence_threshold
+        self._loaded_models: dict[str, YOLO] = {}
+
+        logger.info(
+            f"ZoneDetector initialized: device={self._device}, "
+            f"doc_types={list(self._weights_map.keys())}"
+        )
+
+    def _get_model(self, doc_type: str) -> YOLO:
+        """Загружает модель при первом обращении."""
+        if doc_type not in self._loaded_models:
+            if doc_type not in self._weights_map:
+                raise ValueError(
+                    f"No weights for doc_type='{doc_type}'. "
+                    f"Available: {list(self._weights_map.keys())}"
+                )
+            path = self._weights_map[doc_type]
+            logger.info(f"Loading YOLO model for '{doc_type}': {path}")
+            self._loaded_models[doc_type] = YOLO(path)
+        return self._loaded_models[doc_type]
+
+    @property
+    def supported_doc_types(self) -> list[str]:
+        """Список поддерживаемых типов документов."""
+        return list(self._weights_map.keys())
+
+    def detect(self, image: np.ndarray, doc_type: str) -> list[Detection]:
+        """
+        Обнаруживает зоны на изображении документа.
+
+        Args:
+            image: BGR изображение (кроп одного документа).
+            doc_type: тип документа.
+
+        Returns:
+            Список Detection.
+        """
+        model = self._get_model(doc_type)
+        results = model(image, device=self._device, verbose=False)
+
+        detections = []
+        if results[0].obb is None:
+            return detections
+
+        for det in results[0].obb:
+            confidence = float(det.conf.cpu())
+            if confidence < self._confidence_threshold:
+                continue
+
+            zone_id = int(det.cls.cpu())
+            zone_name = model.names[zone_id]
+            obb_points = det.xyxyxyxy.cpu().numpy().flatten()
+
+            detections.append(Detection(
+                zone_name=zone_name,
+                obb_points=obb_points,
+                confidence=confidence,
+            ))
+
+        return detections

docreader_ocr-0.2.0/src/docreader/factory.py

@@ -0,0 +1,129 @@
+"""
+Фабричные функции для создания компонентов со стандартными весами.
+
+Использование:
+    from docreader import create_classifier, create_detector, create_ocr
+
+    clf = create_classifier()
+    clf = create_classifier(confidence_threshold=0.5) # Переопределение
+
+    det = create_detector()
+    det = create_detector(device="cuda")
+
+    ocr = create_ocr()
+    ocr = create_ocr(lang=["en", "ru"])
+"""
+
+from docreader.config import PipelineConfig
+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
+
+
+def create_classifier(
+    config: PipelineConfig | None = None,
+    **kwargs,
+) -> DocClassifier:
+    """
+    Создаёт классификатор документов со стандартными весами.
+
+    Веса скачиваются автоматически при первом вызове.
+
+    Args:
+        config: конфигурация (если None — используется дефолтная).
+        **kwargs: переопределение параметров DocClassifier
+            (weights_path, device, confidence_threshold).
+
+    Returns:
+        Готовый к работе DocClassifier.
+
+    Примеры:
+        clf = create_classifier()
+        clf = create_classifier(confidence_threshold=0.5)
+        clf = create_classifier(device="cuda")
+    """
+    cfg = config or PipelineConfig()
+
+    defaults = {
+        "weights_path": str(ensure_model(cfg.classifier_weights)),
+        "device": cfg.resolve_device(),
+        "confidence_threshold": cfg.classifier_confidence
+    }
+    defaults.update(kwargs)
+    return DocClassifier(**defaults)
+
+
+def create_detector(
+    config: PipelineConfig | None = None,
+    **kwargs,
+) -> ZoneDetector:
+    """
+    Создаёт детектор зон документов со стандартными весами.
+
+    Args:
+        config: конфигурация (если None — используется дефолтная).
+        **kwargs: переопределение параметров ZoneDetector
+            (weights_map, device, confidence_threshold).
+
+    Returns:
+        Готовый к работе ZoneDetector.
+
+    Примеры:
+        det = create_detector()
+        det = create_detector(device="cuda")
+        det = create_detector(confidence_threshold=0.1)
+    """
+    cfg = config or PipelineConfig()
+
+    weights_map = {
+        doc_type: str(ensure_model(filename))
+        for doc_type, filename in cfg.detector_weights.items()
+    }
+
+    defaults = {
+        "weights_map": weights_map,
+        "device": cfg.resolve_device(),
+        "confidence_threshold": cfg.detector_confidence,
+    }
+    defaults.update(kwargs)
+    return ZoneDetector(**defaults)
+
+
+def create_ocr(
+    config: PipelineConfig | None = None,
+    **kwargs,
+) -> TextRecognizer:
+    """
+    Создаёт OCR-движок со стандартными моделями.
+
+    Args:
+        config: конфигурация (если None — используется дефолтная).
+        **kwargs: переопределение параметров TextRecognizer
+            (lang, gpu, model_storage_directory, и т.д.).
+
+    Returns:
+        Готовый к работе TextRecognizer.
+
+    Примеры:
+        ocr = create_ocr()
+        ocr = create_ocr(lang=["en", "ru"])
+        ocr = create_ocr(gpu=False)
+    """
+    cfg = config or PipelineConfig()
+    easyocr_dir = ensure_model(cfg.ocr_model_archive)
+
+    defaults = {
+        "lang": cfg.ocr_lang,
+        "gpu": cfg.resolve_device != "cpu",
+        "model_storage_directory": str(
+            easyocr_dir / cfg.ocr_model_subdir
+        ),
+        "user_network_directory": str(
+            easyocr_dir / cfg.ocr_network_subdir
+        ),
+        "recog_network": cfg.ocr_recog_network,
+        "download_enabled": cfg.ocr_download_enabled,
+    }
+    defaults.update(kwargs)
+    return TextRecognizer(**defaults)

{docreader_ocr-0.1.2 → docreader_ocr-0.2.0}/src/docreader/hub.py

@@ -17,14 +17,18 @@ from tqdm import tqdm
 
 logger = logging.getLogger(__name__)
 
+_BASE_URL_CLASSIFIER = "https://github.com/mishanyacorleone/docreader/releases/download/v0.2.0"
 _BASE_URL = "https://github.com/mishanyacorleone/docreader/releases/download/v0.1.0"
 
 MODEL_REGISTRY: dict[str, dict] = {
-    "best_doc_classifier.pth": {
-        "url": f"{_BASE_URL}/best_doc_classifier.pth",
-        "sha256": "6d56f45bd33f5296f40bbf32c67a46c01914a3ac7a3dcbbf9aa9a0b8402b59c4",
-        "size_mb": 8.75,
+    # === Классификатор документов (YOLO OBB) ===
+    "doc_classifier.pt": {
+        "url": f"{_BASE_URL_CLASSIFIER}/doc_classifier.pt",
+        "sha256": "b1af689fe58849474a6a5cf879458fcba6d017233ca1bd54b5d83098cd9387f5",
+        "size_mb": 5.49,
     },
+
+    # === Детекторы зон ===
     "passport.pt": {
         "url": f"{_BASE_URL}/passport.pt",
         "sha256": "bebe46bcd4270442c1e14e9b5a403c9f59212d92ed8181af1326f9f80bc0f0c0",
@@ -45,6 +49,8 @@ MODEL_REGISTRY: dict[str, dict] = {
         "sha256": "84775a6ff1ababb3f8e31a8aa768717cf9d65d8b84df9c0cd48eb7bdaf680218",
         "size_mb": 5.82,
     },
+
+    # === EasyOCR ===
     "easyocr_custom.tar.gz": {
         "url": f"{_BASE_URL}/easyocr_custom.tar.gz",
         "sha256": "832ce5a7f3a1086d81beb1c991347e3f545a425646bc87f3f576ae06fecd2420",
@@ -179,4 +185,46 @@ def ensure_all_models(cache_dir: Optional[Path] = None) -> Path:
     cache = cache_dir or get_cache_dir()
     for filename in MODEL_REGISTRY:
         ensure_model(filename, cache)
-    return cache
+    return cache
+
+
+def get_model_paths() -> dict[str, Path]:
+    """
+    Возвращает словарь {имя_модели: полный_путь}
+    для всех зарегистрированных моделей.
+    """
+    cache = get_cache_dir()
+    paths = {}
+
+    for filename, meta in MODEL_REGISTRY.items():
+        if "extract_to" in meta:
+            paths[filename] = cache / meta["extract_to"]
+        else:
+            paths[filename] = cache / filename
+
+    return paths
+
+
+def get_model_status() -> dict[str, dict]:
+    """
+    Показывает статус всех моделей: путь, скачана ли, размер.
+    """
+    cache = get_cache_dir()
+    status = {}
+
+    for filename, meta in MODEL_REGISTRY.items():
+        if "extract_to" in meta:
+            path = cache / meta["extract_to"]
+            exists = path.exists() and any(path.iterdir())
+        else:
+            path = cache / filename
+            exists = path.exists()
+
+        status[filename] = {
+            "path": str(path),
+            "downloaded": exists,
+            "size_mb": meta.get("size_mb", "?"),
+            "url": meta["url"],
+        }
+
+    return status

docreader_ocr-0.2.0/src/docreader/ocr/__init__.py

@@ -0,0 +1,4 @@
+from docreader.ocr.base import BaseOcrEngine, OcrResult
+from docreader.ocr.easyocr_engine import TextRecognizer
+
+__all__ = ["BaseOcrEngine", "OcrResult", "TextRecognizer"]