sinapsis-deepseek-ocr 0.1.0__py3-none-any.whl

sinapsis_deepseek_ocr/helpers/bbox_utils.py

@@ -0,0 +1,18 @@
+def denormalize_bbox(coords: tuple[int, int, int, int], width: int, height: int) -> tuple[int, int, int, int]:
+    """Convert normalized 0-999 coordinates to pixel coordinates.
+
+    Args:
+        coords: (x1, y1, x2, y2) in 0-999 range.
+        width: Image width in pixels.
+        height: Image height in pixels.
+
+    Returns:
+        (x1, y1, x2, y2) in pixel coordinates.
+    """
+    x1, y1, x2, y2 = coords
+    return (
+        round(x1 * width / 999),
+        round(y1 * height / 999),
+        round(x2 * width / 999),
+        round(y2 * height / 999),
+    )

sinapsis_deepseek_ocr/helpers/grounding_parser.py

@@ -0,0 +1,38 @@
+"""Parser for DeepSeek-OCR grounding output."""
+
+import re
+
+from sinapsis_deepseek_ocr.helpers.schemas import GroundingResult
+
+GROUNDING_PATTERN = r"(<\|ref\|>(.*?)<\|/ref\|><\|det\|>(.*?)<\|/det\|>)"
+
+
+def _parse_coordinates(coords_str: str) -> list[tuple[int, int, int, int]] | None:
+    """Parse coordinates string, returning None on failure."""
+    try:
+        coords_list = eval(coords_str)
+        if isinstance(coords_list, list):
+            return [tuple(c) for c in coords_list]
+        return None
+    except (SyntaxError, ValueError):
+        return None
+
+
+def parse_grounding_output(text: str) -> list[GroundingResult]:
+    """Parse grounding tags from model output.
+
+    Args:
+        text: Raw model output containing grounding tags.
+
+    Returns:
+        List of GroundingResult with parsed grounding data.
+    """
+    matches = re.findall(GROUNDING_PATTERN, text, re.DOTALL)
+    results: list[GroundingResult] = []
+
+    for _full_match, label, coords_str in matches:
+        coordinates = _parse_coordinates(coords_str)
+        if coordinates is not None:
+            results.append(GroundingResult(label=label, coordinates=coordinates))
+
+    return results

sinapsis_deepseek_ocr/helpers/mode_registry.py

@@ -0,0 +1,33 @@
+from pydantic import BaseModel
+
+
+class DeepSeekOCRModeConfig(BaseModel):
+    """Configuration for a DeepSeek OCR inference mode.
+
+    Attributes:
+        base_size: The base resolution for image processing.
+        image_size: The target image size for inference.
+        crop_mode: Whether to use crop mode for large images.
+    """
+
+    base_size: int
+    image_size: int
+    crop_mode: bool
+
+
+class DeepSeekOCRModeRegistry:
+    """Registry of predefined DeepSeek OCR mode configurations.
+
+    Attributes:
+        TINY: Configuration for tiny mode (512x512, no crop).
+        SMALL: Configuration for small mode (640x640, no crop).
+        GUNDAM: Configuration for gundam mode (1024 base, 640 image, with crop).
+        BASE: Configuration for base mode (1024x1024, no crop).
+        LARGE: Configuration for large mode (1280x1280, no crop).
+    """
+
+    TINY = DeepSeekOCRModeConfig(base_size=512, image_size=512, crop_mode=False)
+    SMALL = DeepSeekOCRModeConfig(base_size=640, image_size=640, crop_mode=False)
+    GUNDAM = DeepSeekOCRModeConfig(base_size=1024, image_size=640, crop_mode=True)
+    BASE = DeepSeekOCRModeConfig(base_size=1024, image_size=1024, crop_mode=False)
+    LARGE = DeepSeekOCRModeConfig(base_size=1280, image_size=1280, crop_mode=False)

sinapsis_deepseek_ocr/helpers/schemas.py

@@ -0,0 +1,59 @@
+from typing import Literal
+
+import torch
+from pydantic import BaseModel, ConfigDict, model_validator
+from pydantic.dataclasses import dataclass
+from sinapsis_core.utils.env_var_keys import SINAPSIS_CACHE_DIR
+
+_DTYPE_MAP = {"float16": torch.float16, "bfloat16": torch.bfloat16}
+
+
+class DeepSeekOCRInitArgs(BaseModel):
+    """Initialization arguments for the DeepSeek OCR model.
+
+    Note: This model requires CUDA. CPU inference is not supported as DeepSeek's
+    infer() method internally requires CUDA tensors.
+
+    Attributes:
+        pretrained_model_name_or_path (str): HuggingFace model identifier or local path.
+        cache_dir (str): Directory to cache downloaded models.
+        torch_dtype (Literal["float16", "bfloat16", "auto"] | torch.dtype): Precision for model weights.
+        attn_implementation (Literal["flash_attention_2", "eager"]: Attention implementation.
+        trust_remote_code (Literal[True]): Whether to trust remote code from HuggingFace.
+        use_safetensors (Literal[True]): Whether to use safetensors format.
+    """
+
+    pretrained_model_name_or_path: str = "deepseek-ai/DeepSeek-OCR"
+    cache_dir: str = SINAPSIS_CACHE_DIR
+    torch_dtype: Literal["float16", "bfloat16", "auto"] | torch.dtype = "auto"
+    attn_implementation: Literal["flash_attention_2", "eager"] = "flash_attention_2"
+    trust_remote_code: Literal[True] = True
+    use_safetensors: Literal[True] = True
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    @model_validator(mode="after")
+    def resolve_torch_dtype(self) -> "DeepSeekOCRInitArgs":
+        """Resolve 'auto' torch_dtype to 'float16' or 'bfloat16' based on availability.
+
+        Returns:
+            DeepSeekOCRInitArgs: The validated instance with resolved torch_dtype.
+        """
+        if self.torch_dtype == "auto":
+            self.torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+        elif isinstance(self.torch_dtype, str):
+            self.torch_dtype = _DTYPE_MAP.get(self.torch_dtype, self.torch_dtype)
+        return self
+
+
+@dataclass
+class GroundingResult:
+    """A single grounding result with label and bounding box coordinates.
+
+    Attributes:
+        label: The text label for this grounding region.
+        coordinates: List of (x1, y1, x2, y2) tuples in normalized 0-999 range.
+    """
+
+    label: str
+    coordinates: list[tuple[int, int, int, int]]

sinapsis_deepseek_ocr/helpers/tags.py

@@ -0,0 +1,19 @@
+from enum import Enum
+
+
+class Tags(Enum):
+    """UI tags for categorizing DeepSeek OCR templates.
+
+    Attributes:
+        DEEPSEEK: Tag for DeepSeek-related templates.
+        IMAGE: Tag for image processing templates.
+        OCR: Tag for optical character recognition templates.
+        TEXT: Tag for text-related templates.
+        TEXT_RECOGNITION: Tag for text recognition templates.
+    """
+
+    DEEPSEEK = "deepseek"
+    IMAGE = "image"
+    OCR = "optical_character_recognition"
+    TEXT = "text"
+    TEXT_RECOGNITION = "text_recognition"

sinapsis_deepseek_ocr/templates/__init__.py

@@ -0,0 +1,18 @@
+import importlib
+from collections.abc import Callable
+
+_root_lib_path = "sinapsis_deepseek_ocr.templates"
+
+_template_lookup = {
+    "DeepSeekOCRInference": f"{_root_lib_path}.deepseek_ocr_inference",
+}
+
+
+def __getattr__(name: str) -> Callable:
+    if name in _template_lookup:
+        module = importlib.import_module(_template_lookup[name])
+        return getattr(module, name)
+    raise AttributeError(f"template `{name}` not found in {_root_lib_path}")
+
+
+__all__ = list(_template_lookup.keys())

sinapsis_deepseek_ocr/templates/deepseek_ocr_inference.py

@@ -0,0 +1,218 @@
+import gc
+import os
+import tempfile
+from typing import Any, Literal
+
+import cv2
+import torch
+from pydantic import Field
+from sinapsis_core.data_containers.annotations import BoundingBox, ImageAnnotations
+from sinapsis_core.data_containers.data_packet import DataContainer, ImageColor, ImagePacket, TextPacket
+from sinapsis_core.template_base import Template
+from sinapsis_core.template_base.base_models import (
+    OutputTypes,
+    TemplateAttributes,
+    UIPropertiesMetadata,
+)
+from sinapsis_generic_data_tools.helpers.image_color_space_converter_cv import convert_color_space_cv
+from transformers import AutoModel, AutoTokenizer
+
+from sinapsis_deepseek_ocr.helpers.bbox_utils import denormalize_bbox
+from sinapsis_deepseek_ocr.helpers.grounding_parser import parse_grounding_output
+from sinapsis_deepseek_ocr.helpers.mode_registry import DeepSeekOCRModeRegistry
+from sinapsis_deepseek_ocr.helpers.schemas import DeepSeekOCRInitArgs, GroundingResult
+from sinapsis_deepseek_ocr.helpers.tags import Tags
+
+
+class DeepSeekOCRInferenceAttributes(TemplateAttributes):
+    """Attributes for DeepSeek OCR inference.
+
+    Attributes:
+        prompt: The prompt to send to the model (without <image> or <|grounding|> tags).
+        enable_grounding: Whether to enable grounding for bounding box extraction.
+        mode: The inference mode configuration to use.
+        init_args: Initialization arguments for the model.
+    """
+
+    prompt: str = "OCR the image."
+    enable_grounding: bool = False
+    mode: Literal["tiny", "small", "gundam", "base", "large"] = "base"
+    init_args: DeepSeekOCRInitArgs = Field(default_factory=DeepSeekOCRInitArgs)
+
+
+class DeepSeekOCRInference(Template):
+    """Template for DeepSeek OCR inference.
+
+    This template uses the DeepSeek OCR model to extract text from images.
+
+    Attributes:
+        MODE_REGISTRY: Registry of available inference mode configurations.
+        AttributesBaseModel: The Pydantic model for template attributes.
+        UIProperties: UI metadata for the template.
+    """
+
+    MODE_REGISTRY = DeepSeekOCRModeRegistry
+    AttributesBaseModel = DeepSeekOCRInferenceAttributes
+    UIProperties = UIPropertiesMetadata(
+        category="OCR",
+        output_type=OutputTypes.MULTIMODAL,
+        tags=[
+            Tags.DEEPSEEK,
+            Tags.IMAGE,
+            Tags.OCR,
+            Tags.TEXT,
+            Tags.TEXT_RECOGNITION,
+        ],
+    )
+
+    def __init__(self, attributes: TemplateAttributes) -> None:
+        super().__init__(attributes)
+        self.initialize()
+
+    def initialize(self) -> None:
+        """Load the model and tokenizer from pretrained weights."""
+        self.model = self._initialize_model()
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.attributes.init_args.pretrained_model_name_or_path,
+            cache_dir=self.attributes.init_args.cache_dir,
+            trust_remote_code=self.attributes.init_args.trust_remote_code,
+        )
+        self.mode_config = self._get_mode_config()
+        self.full_prompt = self._build_prompt()
+
+    def _initialize_model(self) -> AutoModel:
+        """Initialize and return the model with appropriate configuration.
+
+        Note: Always uses CUDA as DeepSeek's infer() method requires it.
+
+        Returns:
+            AutoModel: The initialized model ready for inference.
+        """
+        model = AutoModel.from_pretrained(**self.attributes.init_args.model_dump())
+        return model.eval().cuda().to(self.attributes.init_args.torch_dtype)
+
+    def _get_mode_config(self) -> dict[str, Any]:
+        """Get the mode configuration for the current inference mode.
+
+        Returns:
+            dict[str, Any]: The mode configuration as a dictionary.
+        """
+        mode_config = getattr(self.MODE_REGISTRY, self.attributes.mode.upper())
+        return mode_config.model_dump()
+
+    def _build_prompt(self) -> str:
+        """Build the full prompt with appropriate tags.
+
+        Returns:
+            Full prompt with <image> and optional <|grounding|> tags.
+        """
+        if self.attributes.enable_grounding:
+            return f"<image>\n<|grounding|>{self.attributes.prompt}"
+        return f"<image>\n{self.attributes.prompt}"
+
+    def infer(self, image_packet: ImagePacket) -> str:
+        """Run OCR inference on an image packet.
+
+        Args:
+            image_packet: The image packet containing the image to process.
+
+        Returns:
+            str: The raw OCR result from the model.
+        """
+        fd, tmp_path = tempfile.mkstemp(suffix=".jpg")
+        os.close(fd)
+
+        try:
+            image_packet = convert_color_space_cv(image_packet=image_packet, desired_color_space=ImageColor.BGR)
+            cv2.imwrite(tmp_path, image_packet.content)
+
+            raw_result = self.model.infer(
+                tokenizer=self.tokenizer,
+                prompt=self.full_prompt,
+                image_file=tmp_path,
+                output_path=self.attributes.init_args.cache_dir,
+                eval_mode=True,
+                save_results=False,
+                **self.mode_config,
+            )
+
+            return raw_result
+
+        finally:
+            if os.path.exists(tmp_path):
+                os.remove(tmp_path)
+
+    @staticmethod
+    def grounding_to_annotations(results: list[GroundingResult], width: int, height: int) -> list[ImageAnnotations]:
+        """Convert GroundingResult list to ImageAnnotations.
+
+        Args:
+            results: List of parsed grounding results.
+            width: Image width in pixels.
+            height: Image height in pixels.
+
+        Returns:
+            List of ImageAnnotations with denormalized bounding boxes.
+        """
+        annotations: list[ImageAnnotations] = []
+
+        for result in results:
+            for coords in result.coordinates:
+                x1, y1, x2, y2 = denormalize_bbox(coords, width, height)
+                bbox = BoundingBox(x=x1, y=y1, w=x2 - x1, h=y2 - y1)
+                ann = ImageAnnotations(
+                    label_str=result.label,
+                    bbox=bbox,
+                    text=result.label,
+                )
+                annotations.append(ann)
+
+        return annotations
+
+    @staticmethod
+    def clear_memory() -> None:
+        """Clear memory to free up resources.
+
+        Performs garbage collection and clears GPU memory if available.
+        """
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    def reset_state(self, template_name: str | None = None) -> None:
+        """Reset the template state by reinitializing the model.
+
+        Args:
+            template_name: Optional template name (unused, for interface compatibility).
+        """
+        _ = template_name
+
+        if hasattr(self, "model"):
+            del self.model
+
+        self.clear_memory()
+        self.initialize()
+        self.logger.info(f"Reset template instance `{self.instance_name}`")
+
+    def execute(self, container: DataContainer) -> DataContainer:
+        """Execute OCR on all images in the container.
+
+        Args:
+            container: The data container with images to process.
+
+        Returns:
+            DataContainer: The container with annotations or text packets.
+        """
+        for image_packet in container.images:
+            raw_result = self.infer(image_packet)
+
+            if self.attributes.enable_grounding:
+                grounding_results = parse_grounding_output(raw_result)
+                height, width = image_packet.shape[:2]
+                annotations = self.grounding_to_annotations(grounding_results, width, height)
+                image_packet.annotations.extend(annotations)
+            else:
+                container.texts.append(TextPacket(content=raw_result))
+
+        return container