openocr-mlt-cuda 0.4.0__py3-none-any.whl

openocr_mlt/__init__.py

@@ -0,0 +1,11 @@
+"""Offline, serial, FP32 ScriptMoE recognition."""
+from .api import OpenOCRMLT
+from .config import RecognitionConfig
+from .types import RecognitionResult, DetectionResult, OCRLine, OCRResult
+from .detection import TextDetector, DetectionConfig
+
+__version__ = "0.4.0"
+__all__ = ["OpenOCRMLT", "RecognitionConfig", "RecognitionResult", "DetectionResult", "OCRLine", "OCRResult", "TextDetector", "DetectionConfig"]
+
+from .model_fetch import fetch_models
+__all__.append("fetch_models")

openocr_mlt/__main__.py

@@ -0,0 +1,2 @@
+from .cli import main
+raise SystemExit(main())

openocr_mlt/api.py

@@ -0,0 +1,275 @@
+"""ScriptMoE recognition, optionally composed with pinned ONNX detection."""
+import os
+from dataclasses import asdict
+from PIL import Image
+import numpy as np
+from .assets import resolve_assets
+from .config import MODEL_ID, RecognitionConfig, resolve_device
+from .processing import read_image, resize_max_side
+from .types import RecognitionResult, DetectionResult, OCRLine, OCRResult
+from .geometry import crop_quad
+
+class _ConfiguredOCR:
+    def __init__(self, *, model_dir=None, weights_path=None, character_dict_path=None,
+                 rec_backend="onnx", rec_device=None, rec_batch_size=1,
+                 offline=True, preset="default", rec_threshold=None, max_side=None,
+                 det_backend=None, det_model_dir=None,
+                 det_model_id="ppocrv6-medium-det-bce-v1", det_device=None,
+                 det_config=None, page_max_side=0, device=None, threads=4):
+        if rec_backend not in ("torch","onnx"):
+            raise ValueError("rec_backend must be torch or onnx")
+        if isinstance(threads,bool) or not isinstance(threads,int) or threads<1:
+            raise ValueError("threads must be a positive integer")
+        common = device if device is not None else "cpu"
+        rec_device = common if rec_device is None else rec_device
+        det_device = common if det_device is None else det_device
+        if det_backend is None and det_model_dir is not None:
+            det_backend = "onnx"
+        self.backend=rec_backend
+        self._closed=False
+        if rec_batch_size != 1:
+            raise ValueError("Only rec_batch_size=1 is supported; lists are processed serially")
+        if offline is not True:
+            raise ValueError("Only offline=True is supported; no online downloader is implemented")
+        if det_backend not in (None,"onnx"):
+            raise ValueError("det_backend must be None or 'onnx'")
+        if det_backend is None and det_model_dir is not None:
+            raise ValueError("det_model_dir requires det_backend='onnx'")
+        if det_backend=="onnx" and det_model_dir is None:
+            raise ValueError("det_backend='onnx' requires det_model_dir")
+        if isinstance(page_max_side,bool) or not isinstance(page_max_side,int) or page_max_side<0:
+            raise ValueError("page_max_side must be an integer >=0")
+        self.config = RecognitionConfig.from_preset(preset, rec_threshold=rec_threshold, max_side=max_side)
+        self.page_max_side=page_max_side
+        self._detector=None
+        if rec_backend=="torch":
+            from .recognition.torch_backend import TorchRecognizer
+            self.device=resolve_device(rec_device)
+            weights,dictionary=resolve_assets(model_dir,weights_path,character_dict_path)
+            self._recognizer=TorchRecognizer(weights,dictionary,self.device,self.config)
+            self.model_id=MODEL_ID
+        else:
+            if weights_path is not None:
+                raise ValueError("ONNX requires model_dir, not a safetensors weights_path")
+            from .recognition.onnx_backend import OnnxRecognizer
+            self._recognizer=OnnxRecognizer(model_dir,character_dict_path,rec_device,self.config,threads)
+            self.device=self._recognizer.device
+            self.model_id=self._recognizer.model_id
+        if det_backend=="onnx" and det_config is None:
+            from .detection import DetectionConfig
+            det_config=DetectionConfig(intra_threads=threads)
+        if det_backend=="onnx":
+            from .detection import TextDetector
+            self._detector=TextDetector(det_model_dir,model_id=det_model_id,device=det_device,config=det_config,offline=offline)
+
+    @property
+    def info(self):
+        info={"model_id":self.model_id,"backend":self.backend,"device":str(self.device),
+              "precision":"float32","effective_batch_size":1,"offline":True,"config":asdict(self.config)}
+        if self._detector is not None:
+            info["detector"]=self._detector.info
+            info["page_max_side"]=self.page_max_side
+        if self.backend=="onnx": info["recognizer"]=self._recognizer.info
+        return info
+
+    def _recognize_loaded(self,image,*,resize):
+        if self._closed: raise RuntimeError("OpenOCRMLT is closed")
+        height,width=image.shape[:2]
+        if resize:
+            image=resize_max_side(image,self.config.max_side)
+        raw,score,termination,steps=self._recognizer.recognize(image)
+        accepted=score>=self.config.rec_threshold
+        status="empty" if not raw else ("ok" if accepted else "filtered")
+        return RecognitionResult(raw if accepted else "",raw,score,status,termination,steps,
+                                 (width,height),self.model_id,self.backend,str(self.device))
+
+    def recognize(self,images,*,color_order="BGR"):
+        """Unchanged rec-only semantics, independent of any configured detector."""
+        if isinstance(images,(str,os.PathLike,Image.Image,np.ndarray)):
+            images=[images]
+        if not isinstance(images,(list,tuple)):
+            raise TypeError("images must be an image or a list/tuple of images")
+        loaded=[read_image(image,color_order) for image in images]
+        return [self._recognize_loaded(image,resize=True) for image in loaded]
+
+    def _detect_page(self,image):
+        if self._detector is None:
+            raise RuntimeError("Configure det_backend='onnx' and det_model_dir before detect/ocr")
+        height,width=image.shape[:2]
+        working=resize_max_side(image,self.page_max_side)
+        detections=self._detector._detect_bgr(working)
+        sy,sx=height/working.shape[0],width/working.shape[1]
+        mapped=[DetectionResult(tuple((min(max(x*sx,0.0),float(width)),min(max(y*sy,0.0),float(height))) for x,y in d.points),
+                                d.score,d.source_id,d.model_id,d.backend,d.device) for d in detections]
+        return working,detections,mapped
+
+    def detect(self,image,*,color_order="BGR"):
+        _,_,mapped=self._detect_page(read_image(image,color_order))
+        return mapped
+
+    def ocr(self,image,*,color_order="BGR"):
+        bgr=read_image(image,color_order)
+        height,width=bgr.shape[:2]
+        working,detections,mapped=self._detect_page(bgr)
+        items=[]
+        for detection,original in zip(detections,mapped):
+            crop,error=crop_quad(working,detection.points)
+            if error is not None:
+                items.append(OCRLine(original,None,error))
+            else:
+                # Legacy whole-page pipeline resizes the page, not every crop.
+                result=self._recognize_loaded(crop,resize=False)
+                items.append(OCRLine(original,result,None))
+        return OCRResult((width,height),tuple(items),self.info)
+
+    def __call__(self,image,*,color_order="BGR"):
+        """One crop without a detector; one page with an explicitly configured detector."""
+        if getattr(self,"_detector",None) is not None:
+            return self.ocr(image,color_order=color_order)
+        return self.recognize([image],color_order=color_order)[0]
+
+    def close(self):
+        if self._detector is not None:
+            self._detector.close()
+        if self._recognizer is not None:
+            if hasattr(self._recognizer,"close"):self._recognizer.close()
+            elif hasattr(self._recognizer,"model"):self._recognizer.model=None
+        self._closed=True
+
+    def __enter__(self):
+        if self._closed:raise RuntimeError("OpenOCRMLT is closed")
+        return self
+
+    def __exit__(self,*args):
+        self.close()
+
+
+class OpenOCRMLT(_ConfiguredOCR):
+    """Use an explicit model configuration or select a task with lazy setup."""
+
+    def __init__(
+        self, *, task=None, device=None, offline=None, cache_dir=None,
+        model_source=None, model_dir=None, det_model_dir=None, **options
+    ):
+        self._task = task
+        if task is None:
+            if cache_dir is not None or model_source is not None:
+                raise ValueError("cache_dir and model_source require a task")
+            super().__init__(
+                device=device, offline=True if offline is None else offline,
+                model_dir=model_dir, det_model_dir=det_model_dir, **options
+            )
+            return
+
+        if task not in ("ocr", "det", "rec"):
+            raise ValueError("task must be ocr, det, or rec")
+        if offline is not None and not isinstance(offline, bool):
+            raise TypeError("offline must be a bool")
+        if task == "det":
+            unsupported = set(options) - {"det_model_id", "det_config", "threads", "det_device"}
+            if unsupported or model_dir is not None:
+                raise ValueError("The detection task accepts detector configuration only")
+        elif task == "rec" and (det_model_dir is not None or options.get("det_backend") is not None):
+            raise ValueError("The recognition task accepts recognizer configuration only")
+        if options.get("rec_backend", "onnx") != "onnx":
+            raise ValueError("Task mode uses ONNX; provide explicit model paths for Torch")
+
+        self._requested_device = device if device is not None else "auto"
+        self._offline = bool(offline)
+        self._cache_dir = cache_dir
+        self._model_source = model_source
+        self._model_dir = model_dir
+        self._det_model_dir = det_model_dir
+        self._options = options
+        self._engine = None
+        self._closed = False
+
+    def _initialize(self):
+        from .model_fetch import fetch_models
+        from .onnx_runtime import _CUDAUnavailableError, resolve_ort_device
+
+        model_dir = self._model_dir
+        det_model_dir = self._det_model_dir
+        if (self._task != "det" and model_dir is None) or (
+            self._task != "rec" and det_model_dir is None
+        ):
+            assets = fetch_models(
+                source=self._model_source, cache_dir=self._cache_dir, offline=self._offline
+            )
+            if model_dir is None:
+                model_dir = assets["recognition_dir"]
+            if det_model_dir is None:
+                det_model_dir = assets["detector_dir"]
+
+        def create(device):
+            if self._task == "det":
+                from .detection import DetectionConfig, TextDetector
+
+                config = self._options.get("det_config")
+                if config is None:
+                    config = DetectionConfig(intra_threads=self._options.get("threads", 4))
+                return TextDetector(
+                    det_model_dir, config=config,
+                    model_id=self._options.get("det_model_id", "ppocrv6-medium-det-bce-v1"),
+                    device=self._options.get("det_device") or device,
+                )
+            return _ConfiguredOCR(
+                model_dir=model_dir,
+                det_model_dir=det_model_dir if self._task == "ocr" else None,
+                device=device, **self._options
+            )
+
+        requested = self._requested_device
+        device = resolve_ort_device(requested)
+        try:
+            return create(device)
+        except _CUDAUnavailableError:
+            if requested != "auto" or device == "cpu" or any(
+                self._options.get(key) for key in ("det_device", "rec_device")
+            ):
+                raise
+            return create("cpu")
+
+    def _task_engine(self, task):
+        if self._closed:
+            raise RuntimeError("OpenOCRMLT is closed")
+        if self._task != task:
+            raise ValueError(f"This instance was configured for task={self._task!r}")
+        if self._engine is None:
+            self._engine = self._initialize()
+        return self._engine
+
+    def __call__(self, image=None, *, image_path=None, color_order="BGR"):
+        if (image is None) == (image_path is None):
+            raise ValueError("Provide exactly one image or image_path")
+        image = image if image_path is None else image_path
+        if self._task is None:
+            return super().__call__(image, color_order=color_order)
+        method = {"ocr": self.ocr, "det": self.detect, "rec": self.recognize}[self._task]
+        return method(image, color_order=color_order)
+
+    def recognize(self, images, *, color_order="BGR"):
+        engine = super() if self._task is None else self._task_engine("rec")
+        return engine.recognize(images, color_order=color_order)
+
+    def detect(self, image, *, color_order="BGR"):
+        engine = super() if self._task is None else self._task_engine("det")
+        return engine.detect(image, color_order=color_order)
+
+    def ocr(self, image, *, color_order="BGR"):
+        engine = super() if self._task is None else self._task_engine("ocr")
+        return engine.ocr(image, color_order=color_order)
+
+    @property
+    def info(self):
+        if self._task is None:
+            return super().info
+        info = {} if self._engine is None else self._engine.info
+        return {**info, "task": self._task, "requested_device": self._requested_device}
+
+    def close(self):
+        if self._task is None:
+            return super().close()
+        if self._engine is not None:
+            self._engine.close()
+        self._closed = True

openocr_mlt/assets/THIRD_PARTY_NOTICES.txt

@@ -0,0 +1,26 @@
+# Third-Party Model Sources and Licensing
+
+## ScriptMoE Recognition
+
+Recognition models are derived from [ScriptMoE](https://github.com/YesianRohn/ScriptMoE) at upstream revision `5dde99867aaee6b4af3e75a3020b225eab17ac53`. ONNX conversion preserves the character dictionary and autoregressive decoding rules and uses KV caching and optimized MoE execution.
+
+| Artifact | SHA256 |
+|---|---|
+| Original recognition weights | 60c74a924404b652c3a9550ec370c74d78850151785341729bdac268c014650d |
+| Character dictionary | 7b4eb2ca5fd6c9b6c6ecf423f69d35d8da48088697303e53624af196ee9ab6c0 |
+
+## PP-OCRv6 Medium Detection
+
+The detector is distributed by Paddle as [PP-OCRv6_medium_det_onnx.tar](https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/tmp/PP-OCRv6_medium_det_onnx.tar).
+
+| Property | Value |
+|---|---|
+| File | inference.onnx |
+| Size | 62,026,563 bytes |
+| SHA256 | a0d14865d3a297970899b5c246f0e47c43879bb406c3a1fa4e505aba84b134e4 |
+
+The accompanying Python package adapts detection preprocessing and postprocessing from [ppocrv6_onnx](https://github.com/AIwork4me/ppocrv6_onnx), revision `d638eb0bb5d8222522abcc1bae4394cbeeb636f0`, under its MIT license. The package includes the license text for that code.
+
+## Licensing Status
+
+A unified redistribution license has not been specified for this combined model bundle. The MIT license of the detection adapter applies to that code; model weights have separate licensing requirements.

openocr_mlt/assets/detector-models.json

@@ -0,0 +1,37 @@
+{
+  "schema_version": 1,
+  "models": {
+    "ppocrv6-tiny-det-bce-v1": {
+      "filename": "inference.onnx",
+      "size": 1780621,
+      "sha256": "a56a3430a96a6c691f8bfbbb297208bb9573c3d70083d6382071cbd92d0a152e",
+      "source_url": "https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/tmp/PP-OCRv6_tiny_det_onnx.tar",
+      "source_revision": null,
+      "artifact_etag": "700445c8a6ebcaeac88d4622ddb68dba",
+      "runtime_profile": "ppocrv6-reference-v1",
+      "config_sha256": "626e4e3b45bd2f943d2282114b461ef5804f4620ae54d7ff7ff9fb6365f3a608",
+      "license_status": "official family model card states Apache-2.0; BCE binary redistribution not independently resolved"
+    },
+    "ppocrv5-server-det-hf-v1": {
+      "filename": "inference.onnx",
+      "size": 88116791,
+      "sha256": "10803475a591f7dc623e24670fb5752ec94d39a1f8cf069aac1b6f0ce19cfc85",
+      "source_url": "https://huggingface.co/PaddlePaddle/PP-OCRv5_server_det_onnx",
+      "source_revision": "dcf248c4dcc064c6d030db241255f9e8fbc6733a",
+      "runtime_profile": "ppocrv6-reference-v1",
+      "license_status": "official model card Apache-2.0; no public release in this task"
+    },
+    "ppocrv6-medium-det-bce-v1": {
+      "name": "inference.onnx",
+      "size": 62026563,
+      "sha256": "a0d14865d3a297970899b5c246f0e47c43879bb406c3a1fa4e505aba84b134e4",
+      "filename": "inference.onnx",
+      "source_url": "https://paddle-model-ecology.bj.bcebos.com/paddlex/official_inference_model/paddle3.0.0/tmp/PP-OCRv6_medium_det_onnx.tar",
+      "source_revision": null,
+      "source_etag": "\"33e0718cf119c8fcc545eb249bc69e06\"",
+      "archive_sha256": "46849788192d7882e52440d81cdfa437156f433368411c04b832ffa5433e5418",
+      "runtime_profile": "ppocrv6-reference-v1",
+      "license_status": "Source attribution preserved; no public redistribution performed"
+    }
+  }
+}