asterocr 1.0.0__py3-none-any.whl

asterocr/models/model.py

@@ -0,0 +1,113 @@
+"""
+AsterNet - End-to-End Deep Neural Network for Vietnamese Text Recognition
+Architecture:
+  Input (B, 3, 32, W)
+    └──> High-Frequency Asymmetric ResNet-SVTR Backbone  --> (T = W/4, B, 512)
+    └──> AsterBiLSTMSequence (2-Layer Bidirectional)      --> (T, B, 512)
+    ├──> Main CTC Classifier Head                         --> Log-Probs (T, B, 227)
+    └──> Diacritic Auxiliary Heads (Optional Training):
+          ├──> Tone Classifier (T, B, 6) [None, Sắc, Huyền, Hỏi, Ngã, Nặng]
+          └──> Hat Classifier  (T, B, 5) [None, Mũ ^, Móc ư/ơ, Trăng ă, Gạch đ]
+"""
+
+import torch
+from torch import nn
+from asterocr.models.backbone import AsterResNetBackbone
+from asterocr.models.sequence import AsterBiLSTMSequence
+from asterocr.models.language import AsterLanguageModelHead, GatedLanguageFusion
+from asterocr.vocab import default_vocab
+
+class AsterNet(nn.Module):
+    def __init__(self, num_classes=None, in_channels=3, feat_dim=512, hidden_dim=256, use_ilm=True):
+        super(AsterNet, self).__init__()
+        if num_classes is None:
+            num_classes = len(default_vocab)
+            
+        self.num_classes = num_classes
+        self.use_ilm = use_ilm
+        self.backbone = AsterResNetBackbone(in_channels=in_channels, out_dim=feat_dim)
+        self.sequence = AsterBiLSTMSequence(in_dim=feat_dim, hidden_dim=hidden_dim)
+        seq_dim = hidden_dim * 2 # 512
+        
+        # Implicit Language Modeling Head (ILMH) & Gated Vision-Language Fusion
+        if self.use_ilm:
+            self.language_head = AsterLanguageModelHead(d_model=seq_dim, nhead=4, num_layers=2)
+            self.fusion = GatedLanguageFusion(d_model=seq_dim, init_vision_bias=1.5)
+            # Auxiliary vision projection for multi-task loss supervision
+            self.vision_classifier = nn.Linear(seq_dim, num_classes)
+            
+        # Main Fused CTC projection head
+        self.classifier = nn.Linear(seq_dim, num_classes)
+        self.log_softmax = nn.LogSoftmax(dim=2)
+        
+        # Diacritic auxiliary heads for supervision against diacritic loss/swallowing
+        self.tone_classifier = nn.Linear(seq_dim, 6) # None, Sắc, Huyền, Hỏi, Ngã, Nặng
+        self.hat_classifier = nn.Linear(seq_dim, 5)  # None, Mũ ^, Móc ư/ơ, Trăng ă, Gạch đ
+
+    def forward(self, x, return_aux=False):
+        """
+        Forward pass for CTC Loss training or evaluation.
+        Args:
+            x: Tensor of shape (B, 3, 32, W)
+            return_aux: bool, if True returns (log_probs, tone_logits, hat_logits, [vis_log_probs])
+        Returns:
+            log_probs: Tensor of shape (T, B, num_classes) where T = W / 4
+        """
+        feats = self.backbone(x)          # (T, B, feat_dim)
+        h_vis = self.sequence(feats)      # (T, B, seq_dim)
+        
+        if self.use_ilm:
+            h_lang = self.language_head(h_vis)
+            h_fused, gate = self.fusion(h_vis, h_lang)
+            logits = self.classifier(h_fused)
+            log_probs = self.log_softmax(logits)
+            
+            if return_aux:
+                vis_logits = self.vision_classifier(h_vis)
+                vis_log_probs = self.log_softmax(vis_logits)
+                tone_logits = self.tone_classifier(h_fused)
+                hat_logits = self.hat_classifier(h_fused)
+                return log_probs, tone_logits, hat_logits, vis_log_probs
+        else:
+            logits = self.classifier(h_vis)
+            log_probs = self.log_softmax(logits)
+            if return_aux:
+                tone_logits = self.tone_classifier(h_vis)
+                hat_logits = self.hat_classifier(h_vis)
+                return log_probs, tone_logits, hat_logits
+            
+        return log_probs
+
+    def load_pretrained(self, checkpoint_path, device="cpu"):
+        """Gracefully loads checkpoint weights, supporting partial loads without errors."""
+        ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        state_dict = ckpt.get("state_dict", ckpt.get("model_state_dict", ckpt))
+        missing, unexpected = self.load_state_dict(state_dict, strict=False)
+        print(f"[+] Loaded weights from {checkpoint_path} (Missing keys: {len(missing)}, Unexpected: {len(unexpected)})")
+        return self
+
+    @torch.inference_mode()
+    def predict_tokens(self, x):
+        """Returns argmax token indices (T, B)."""
+        log_probs = self.forward(x, return_aux=False)
+        return torch.argmax(log_probs, dim=2)
+
+if __name__ == "__main__":
+    net = AsterNet(num_classes=227, use_ilm=True)
+    x = torch.randn(4, 3, 32, 280)
+    out, tone, hat, vis = net(x, return_aux=True)
+    print(f"AsterNet input       : {x.shape}")
+    print(f"AsterNet main output : {out.shape} (T={out.shape[0]}, Batch={out.shape[1]}, Classes={out.shape[2]})")
+    print(f"AsterNet tone output : {tone.shape} (T={tone.shape[0]}, Batch={tone.shape[1]}, Tones=6)")
+    print(f"AsterNet hat output  : {hat.shape} (T={hat.shape[0]}, Batch={hat.shape[1]}, Hats=5)")
+    print(f"AsterNet vis output  : {vis.shape} (T={vis.shape[0]}, Batch={vis.shape[1]}, Classes={vis.shape[2]})")
+    total_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
+    print(f"Total Trainable Parameters: {total_params / 1e6:.2f}M")
+    assert out.shape == (70, 4, 227)
+    
+    # Test loading previous checkpoint without breaking
+    import os
+    ckpt_p = "models/asternet_last.pth"
+    if os.path.exists(ckpt_p):
+        net.load_pretrained(ckpt_p)
+    print("[+] Model verification passed!")

asterocr/models/sequence.py

@@ -0,0 +1,34 @@
+"""
+AsterOCR Sequence Modeling Layer
+Uses Bidirectional LSTM to capture bidirectional contextual dependencies
+between consecutive Vietnamese characters and diacritical marks.
+"""
+
+import torch
+from torch import nn
+
+class AsterBiLSTMSequence(nn.Module):
+    def __init__(self, in_dim=512, hidden_dim=256, num_layers=2, dropout=0.1):
+        super(AsterBiLSTMSequence, self).__init__()
+        self.lstm = nn.LSTM(
+            input_size=in_dim,
+            hidden_size=hidden_dim,
+            num_layers=num_layers,
+            bidirectional=True,
+            dropout=dropout if num_layers > 1 else 0.0,
+            batch_first=False
+        )
+        self.out_dim = hidden_dim * 2
+
+    def forward(self, x):
+        # x shape: (T, B, in_dim)
+        recurrent, _ = self.lstm(x)
+        # recurrent shape: (T, B, hidden_dim * 2)
+        return recurrent
+
+if __name__ == "__main__":
+    seq = AsterBiLSTMSequence(in_dim=512, hidden_dim=256)
+    dummy_feat = torch.randn(64, 2, 512) # T=64, B=2, D=512
+    out = seq(dummy_feat)
+    print("Sequence input shape:", dummy_feat.shape)
+    print("Sequence output shape:", out.shape)

asterocr/pipeline.py

@@ -0,0 +1,179 @@
+"""
+AsterOCR Main High-Level Engine
+Universal Vietnamese OCR Predictor with ONNX Runtime & PyTorch Backends.
+Features automatic Hugging Face model weight acquisition, caching, and CLI support.
+"""
+
+import os
+import sys
+import json
+import numpy as np
+from PIL import Image
+
+try:
+    from huggingface_hub import hf_hub_download
+except ImportError:
+    hf_hub_download = None
+
+from asterocr.vocab import default_vocab
+
+DEFAULT_HF_REPO = "asterxlabs/asternet-vietnamese-ocr"
+CACHE_DIR = os.path.expanduser(os.path.join("~", ".cache", "asterocr"))
+
+class AsterOCR:
+    """
+    Main user-facing OCR recognition engine for AsterOCR.
+    Usage:
+        >>> from asterocr import AsterOCR
+        >>> ocr = AsterOCR()
+        >>> text = ocr.predict("invoice.png")
+        >>> print(text)
+    """
+    def __init__(self, model_path=None, backend="onnx", device="cpu", enable_telemetry=True):
+        self.backend = backend.lower()
+        self.device = device
+        self.vocab = default_vocab
+        self.enable_telemetry = enable_telemetry
+        
+        self.model_path = self._resolve_model_path(model_path)
+        
+        if self.backend == "onnx":
+            self._init_onnx()
+        elif self.backend == "torch":
+            self._init_torch()
+        else:
+            raise ValueError(f"Unsupported backend '{backend}'. Choose 'onnx' or 'torch'.")
+
+    def _resolve_model_path(self, model_path):
+        if model_path and os.path.exists(model_path):
+            return model_path
+            
+        filename = "asternet.onnx" if self.backend == "onnx" else "asternet_best.pth"
+        
+        # 1. Local workspace candidates
+        candidates = [
+            os.path.join("models", filename),
+            os.path.join("hf_model", filename),
+            os.path.join(CACHE_DIR, filename),
+            os.path.join(os.path.dirname(__file__), "..", "models", filename),
+        ]
+        for c in candidates:
+            if os.path.exists(c):
+                return c
+
+        # 2. Auto-fetch from Hugging Face Hub
+        os.makedirs(CACHE_DIR, exist_ok=True)
+        print(f"[*] Đang tải mô hình AsterOCR ({filename}) từ Hugging Face Hub ({DEFAULT_HF_REPO})...")
+        if hf_hub_download is None:
+            raise ImportError("Vui lòng cài đặt 'huggingface_hub' để tự động tải model: pip install huggingface_hub")
+            
+        cached_file = hf_hub_download(
+            repo_id=DEFAULT_HF_REPO,
+            filename=filename,
+            cache_dir=CACHE_DIR
+        )
+        print(f"[+] Đã tải mô hình thành công về: {cached_file}")
+        return cached_file
+
+    def _init_onnx(self):
+        try:
+            import onnxruntime as ort
+        except ImportError:
+            raise ImportError("Vui lòng cài đặt 'onnxruntime' để chạy backend ONNX: pip install onnxruntime")
+            
+        providers = ["CPUExecutionProvider"]
+        if self.device == "cuda" and "CUDAExecutionProvider" in ort.get_available_providers():
+            providers.insert(0, "CUDAExecutionProvider")
+            
+        self.session = ort.InferenceSession(self.model_path, providers=providers)
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_name = self.session.get_outputs()[0].name
+
+    def _init_torch(self):
+        try:
+            import torch
+            from asterocr.models.model import AsterNet
+        except ImportError:
+            raise ImportError("Vui lòng cài đặt 'torch' để chạy backend PyTorch: pip install torch")
+            
+        self.torch_device = torch.device(self.device if torch.cuda.is_available() or self.device == "cpu" else "cpu")
+        self.model = AsterNet(num_classes=len(self.vocab)).to(self.torch_device)
+        self.model.load_pretrained(self.model_path, device=self.torch_device)
+        self.model.eval()
+
+    def preprocess_image(self, image_input):
+        """Converts input image to normalized numpy tensor of shape (1, 3, 32, W)."""
+        if isinstance(image_input, str):
+            image = Image.open(image_input)
+        elif isinstance(image_input, bytes):
+            import io
+            image = Image.open(io.BytesIO(image_input))
+        elif isinstance(image_input, Image.Image):
+            image = image_input
+        elif isinstance(image_input, np.ndarray):
+            image = Image.fromarray(image_input)
+        else:
+            raise TypeError(f"Unsupported image input type: {type(image_input)}")
+
+        image = image.convert("RGB")
+        w, h = image.size
+        target_h = 32
+        target_w = max(32, int(w * target_h / h))
+        target_w = (target_w + 3) // 4 * 4  # multiple of 4
+
+        resized = image.resize((target_w, target_h), Image.Resampling.BILINEAR)
+        arr = np.array(resized, dtype=np.float32) / 255.0
+        # Standard mean=0.5, std=0.5 normalization
+        arr = (arr - 0.5) / 0.5
+        # (H, W, C) -> (1, C, H, W)
+        tensor = np.transpose(arr, (2, 0, 1))[np.newaxis, ...]
+        return tensor
+
+    def decode_ctc(self, log_probs_numpy):
+        """
+        Fast Greedy CTC Decoder with Blank removal and run-length collapsing.
+        log_probs_numpy: shape (T, 1, C) or (1, T, C)
+        """
+        if log_probs_numpy.ndim == 3:
+            if log_probs_numpy.shape[1] == 1:
+                # Shape (T, 1, C)
+                best_indices = np.argmax(log_probs_numpy[:, 0, :], axis=-1)
+            else:
+                # Shape (1, T, C)
+                best_indices = np.argmax(log_probs_numpy[0, :, :], axis=-1)
+        else:
+            best_indices = np.argmax(log_probs_numpy, axis=-1)
+
+        blank_idx = self.vocab.blank_idx
+        prev_idx = -1
+        decoded_chars = []
+        for idx in best_indices:
+            idx = int(idx)
+            if idx != prev_idx:
+                if idx != blank_idx:
+                    decoded_chars.append(self.vocab.idx2char.get(idx, ""))
+                prev_idx = idx
+        return "".join(decoded_chars)
+
+    def predict(self, image_input):
+        """
+        Runs inference on an image and returns recognized Vietnamese text string.
+        """
+        tensor = self.preprocess_image(image_input)
+        
+        if self.backend == "onnx":
+            ort_inputs = {self.input_name: tensor}
+            ort_output = self.session.run([self.output_name], ort_inputs)[0]
+            text = self.decode_ctc(ort_output)
+        else:
+            import torch
+            with torch.no_grad():
+                pt_tensor = torch.from_numpy(tensor).to(self.torch_device)
+                pt_output = self.model(pt_tensor)
+                text = self.decode_ctc(pt_output.cpu().numpy())
+                
+        return text
+
+    def predict_batch(self, image_list):
+        """Runs sequential or batched inference on a list of images."""
+        return [self.predict(img) for img in image_list]