fastplateocr-py 0.1.0__py3-none-any.whl

fastplateocr/__init__.py

@@ -0,0 +1,5 @@
+__version__ = '0.1.0'
+
+from .pipeline import FastPlateOCR
+
+__all__ = ['FastPlateOCR']

fastplateocr/det/__init__.py

@@ -0,0 +1 @@
+# det package

fastplateocr/det/c3k2/bottleneck/rep_mixer.py

@@ -0,0 +1,81 @@
+import torch
+import torch.nn as nn
+
+
+class RepMixerBlock(nn.Module):
+    """
+    Macro-architectural Reparameterization (Large Kernel 5x5).
+    Replace two consecutive 3x3 layers with a single 5x5 layer.
+    During training, use 4 branches (5x5, 3x3, 1x1, Identity)
+    to maximize the network's feature extraction capability.
+    """
+    def __init__(self, in_channels, out_channels, deploy=False):
+        super().__init__()
+        self.deploy = deploy
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.act = nn.SiLU(inplace=True)
+
+        if deploy:
+            # At deploy time: Fuse into a single 5x5 matrix
+            self.rbr_reparam = nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=2, bias=True)
+        else:
+            # Branch 1: 5x5 (Ultra-wide receptive field)
+            self.rbr_5x5 = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=2, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+            # Branch 2: 3x3 (Learn local details)
+            self.rbr_3x3 = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+            # Branch 3: 1x1 (Learn channel cross-talk)
+            self.rbr_1x1 = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0, bias=False),
+                nn.BatchNorm2d(out_channels)
+            )
+            # Branch 4: Identity (Skip computation)
+            if in_channels == out_channels:
+                self.rbr_identity = nn.BatchNorm2d(in_channels)
+            else:
+                self.rbr_identity = None
+
+    def forward(self, x):
+        if self.deploy:
+            return self.act(self.rbr_reparam(x))
+
+        # Thanks to standard padding (2, 1, 0), all 3 conv branches output Tensors with the same H, W.
+        out = self.rbr_5x5(x) + self.rbr_3x3(x) + self.rbr_1x1(x)
+        if self.rbr_identity is not None:
+            out += self.rbr_identity(x)
+        return self.act(out)
+
+
+def setup_monkey_patch_injection():
+    """
+    Remove the entire old Bottleneck structure (including cv1 3x3 and cv2 3x3).
+    Replace it with a 5x5 RepMixerBlock going straight from input channel (c1) to output channel (c2).
+    """
+    import ultralytics.nn.modules.block as block
+
+    def new_init(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5):
+        # Initialize the base nn.Module class (Discard all old Bottleneck code)
+        nn.Module.__init__(self)
+        self.c1 = c1
+        self.c2 = c2
+        # Shortcut only activated if c1 == c2 (keeping YOLO's rule)
+        self.add = shortcut and c1 == c2
+
+        # ARCHITECTURAL BREAKTHROUGH: Remove the channel compression layer (c_)
+        # Use exactly 1 multi-branch Large Kernel block.
+        self.rep_block = RepMixerBlock(c1, c2, deploy=False)
+
+    def new_forward(self, x):
+        # Straight signal transmission through RepMixerBlock
+        return x + self.rep_block(x) if self.add else self.rep_block(x)
+
+    # Overwrite YOLO architecture
+    block.Bottleneck.__init__ = new_init
+    block.Bottleneck.forward = new_forward
+    print("[INFO] Monkey-Patch successful: Overhauled Bottleneck architecture with multi-branch 5x5 Large Kernel!")

fastplateocr/license_plate_dict.txt

@@ -0,0 +1,37 @@
+0
+1
+2
+3
+4
+5
+6
+7
+8
+9
+A
+B
+C
+D
+Đ
+E
+F
+G
+H
+I
+J
+K
+L
+M
+N
+O
+P
+Q
+R
+S
+T
+U
+V
+W
+X
+Y
+Z

fastplateocr/pipeline.py

@@ -0,0 +1,185 @@
+import os
+import sys
+import torch
+import cv2
+import numpy as np
+import warnings
+warnings.filterwarnings("ignore")
+
+# This ensures that we can import det and rec if running from this file
+__dir__ = os.path.dirname(os.path.abspath(__file__))
+root_dir = os.path.abspath(os.path.join(__dir__, '..'))
+if root_dir not in sys.path:
+    sys.path.insert(0, root_dir)
+
+from ultralytics import YOLO
+from fastplateocr.det.c3k2.bottleneck.rep_mixer import setup_monkey_patch_injection, RepMixerBlock
+from fastplateocr.rec.modeling import build_model
+from fastplateocr.rec.postprocess import build_post_process
+
+# Trick for PyTorch Unpickler: Inject RepMixerBlock into __main__
+import __main__
+if not hasattr(__main__, 'RepMixerBlock'):
+    setattr(__main__, 'RepMixerBlock', RepMixerBlock)
+
+
+def switch_deploy_flag(model):
+    """Recursively convert RepMixer layers in SVTR to deploy mode."""
+    for m in model.modules():
+        if hasattr(m, 'switch_to_deploy'):
+            m.switch_to_deploy()
+    return model
+
+
+def download_from_hf(filename):
+    try:
+        from huggingface_hub import hf_hub_download
+        print(f"[FastPlateOCR] Downloading/Verifying {filename} from HuggingFace...")
+        return hf_hub_download(repo_id="anhone3/FastPlateOCR", filename=filename)
+    except ImportError:
+        raise ImportError("Please install huggingface_hub to auto-download models: pip install huggingface_hub")
+
+
+class FastPlateOCR:
+    def __init__(self,
+                 use_det=True,
+                 use_rec=True,
+                 det_model_path=None,
+                 rec_model_path=None,
+                 device=None):
+
+        self.device = torch.device(device if device else ('cuda:0' if torch.cuda.is_available() else 'cpu'))
+        print(f"[FastPlateOCR] Initializing on {self.device}...")
+
+        self.det_model = None
+        self.rec_model = None
+        self.post_process_class = None
+
+        # 1. Initialize DET (YOLO)
+        if use_det:
+            if det_model_path is None:
+                det_model_path = download_from_hf("yolo26n_rep_mixer/best.pt")
+
+            print(f"[FastPlateOCR] Loading Detection Model: {det_model_path}")
+            setup_monkey_patch_injection()
+            self.det_model = YOLO(det_model_path)
+            self.det_model.to(self.device)
+            self.det_model.fuse()
+
+            # Fuse RepMixerBlock inside YOLO if they exist
+            if hasattr(self.det_model, 'model'):
+                for m in self.det_model.model.modules():
+                    if hasattr(m, 'switch_to_deploy'):
+                        m.switch_to_deploy()
+
+        # 2. Initialize REC (SVTRv2)
+        if use_rec:
+            if rec_model_path is None:
+                rec_model_path = download_from_hf("svtrv2_tiny_efficient_rctc/best.pth")
+
+            print(f"[FastPlateOCR] Loading Recognition Model: {rec_model_path}")
+            checkpoint = torch.load(rec_model_path, map_location='cpu')
+            cfg = checkpoint['config']
+
+            # Monkey patch the dictionary path
+            dict_path = os.path.join(os.path.dirname(__file__), 'license_plate_dict.txt')
+            cfg['Global']['character_dict_path'] = dict_path
+
+            self.post_process_class = build_post_process(cfg['PostProcess'], cfg['Global'])
+            cfg['Architecture']['Decoder']['out_channels'] = self.post_process_class.get_character_num()
+
+            self.rec_model = build_model(cfg['Architecture'])
+            self.rec_model.load_state_dict(checkpoint['state_dict'], strict=True)
+            self.rec_model.to(self.device)
+            self.rec_model.eval()
+            self.rec_model = switch_deploy_flag(self.rec_model)
+
+        if self.det_model or self.rec_model:
+            print("[FastPlateOCR] Models loaded successfully!")
+        else:
+            print("[FastPlateOCR] WARNING: No models loaded. Please provide det_model_path or rec_model_path.")
+
+    def _preprocess_crop(self, img_crop, target_h=32, target_w=128):
+        """Preprocesses cropped image for SVTR."""
+        img = cv2.resize(img_crop, (target_w, target_h))
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        img = img.astype(np.float32) / 255.0
+        img -= 0.5
+        img /= 0.5
+        img = img.transpose((2, 0, 1))
+        img = np.expand_dims(img, axis=0)
+        return torch.from_numpy(img).to(self.device)
+
+    def detect(self, img, conf_thresh=0.25):
+        """
+        Detect license plates in an image.
+        Returns: list of bounding boxes [[x1, y1, x2, y2], ...]
+        """
+        if self.det_model is None:
+            raise ValueError("Detection model not loaded! Initialize with det_model_path.")
+
+        if isinstance(img, str):
+            img = cv2.imread(img)
+
+        det_results = self.det_model(img, verbose=False, conf=conf_thresh)[0]
+        boxes = det_results.boxes.data.cpu().numpy() # [x1, y1, x2, y2, conf, cls]
+
+        results = []
+        for box in boxes:
+            x1, y1, x2, y2, conf, cls = box
+            results.append([int(x1), int(y1), int(x2), int(y2)])
+        return results
+
+    def recognize(self, crop_img):
+        """
+        Recognize text in a cropped license plate image.
+        Returns: tuple (text, confidence_score)
+        """
+        if self.rec_model is None:
+            raise ValueError("Recognition model not loaded! Initialize with rec_model_path.")
+
+        if crop_img.size == 0:
+            return "", 0.0
+
+        tensor = self._preprocess_crop(crop_img)
+        with torch.no_grad():
+            preds = self.rec_model(tensor)
+
+        post_result = self.post_process_class(preds)
+        text, score = post_result[0]
+        return text, float(score)
+
+    def read(self, image_path, conf_thresh=0.25):
+        """
+        End-to-End inference.
+        Returns: list of dicts [{'box': [x1,y1,x2,y2], 'text': '51F1234', 'score': 0.99}]
+        """
+        if self.det_model is None or self.rec_model is None:
+            raise ValueError("End-to-End read() requires BOTH det_model_path and rec_model_path to be loaded.")
+
+        img = cv2.imread(image_path)
+        if img is None:
+            raise ValueError(f"Could not read image: {image_path}")
+
+        boxes = self.detect(img, conf_thresh)
+        final_results = []
+
+        h, w, _ = img.shape
+        for box in boxes:
+            x1, y1, x2, y2 = box
+
+            # Ensure bounds
+            x1_c, y1_c = max(0, x1), max(0, y1)
+            x2_c, y2_c = min(w, x2), min(h, y2)
+
+            crop_img = img[y1_c:y2_c, x1_c:x2_c]
+
+            text, score = self.recognize(crop_img)
+
+            final_results.append({
+                'box': box,
+                'text': text,
+                'score': score
+            })
+
+        return final_results

fastplateocr/rec/losses/__init__.py

@@ -0,0 +1,71 @@
+import copy
+from importlib import import_module
+from torch import nn
+
+name_to_module = {
+    'ABINetLoss': '.abinet_loss',
+    'ARLoss': '.ar_loss',
+    'CDistNetLoss': '.cdistnet_loss',
+    'CELoss': '.ce_loss',
+    'CPPDLoss': '.cppd_loss',
+    'CTCLoss': '.ctc_loss',
+    'IGTRLoss': '.igtr_loss',
+    'LISTERLoss': '.lister_loss',
+    'LPVLoss': '.lpv_loss',
+    'MGPLoss': '.mgp_loss',
+    'PARSeqLoss': '.parseq_loss',
+    'RobustScannerLoss': '.robustscanner_loss',
+    'SEEDLoss': '.seed_loss',
+    'SMTRLoss': '.smtr_loss',
+    'SRNLoss': '.srn_loss',
+    'VisionLANLoss': '.visionlan_loss',
+    'CAMLoss': '.cam_loss',
+    'MDiffLoss': '.mdiff_loss',
+    'UniRecLoss': '.unirec_loss',
+    'CMERLoss': '.cmer_loss',
+}
+
+
+def build_loss(config):
+    config = copy.deepcopy(config)
+    module_name = config.pop('name')
+
+    if module_name in globals():
+        module_class = globals()[module_name]
+    else:
+        assert module_name in name_to_module, Exception(
+            '{} is not supported. The losses in {} are supportes'.format(
+                module_name, list(name_to_module.keys())))
+        module_path = name_to_module[module_name]
+        module = import_module(module_path, package=__package__)
+        module_class = getattr(module, module_name)
+
+    return module_class(**config)
+
+
+class GTCLoss(nn.Module):
+
+    def __init__(self,
+                 gtc_loss,
+                 gtc_weight=1.0,
+                 ctc_weight=1.0,
+                 zero_infinity=True,
+                 **kwargs):
+        super(GTCLoss, self).__init__()
+        # Dynamically build CTCLoss
+        ctc_config = {'name': 'CTCLoss', 'zero_infinity': zero_infinity}
+        self.ctc_loss = build_loss(ctc_config)
+        # Build GTC loss
+        self.gtc_loss = build_loss(gtc_loss)
+        self.gtc_weight = gtc_weight
+        self.ctc_weight = ctc_weight
+
+    def forward(self, predicts, batch):
+        ctc_loss = self.ctc_loss(predicts['ctc_pred'],
+                                 [None] + batch[-2:])['loss']
+        gtc_loss = self.gtc_loss(predicts['gtc_pred'], batch[:-2])['loss']
+        return {
+            'loss': self.ctc_weight * ctc_loss + self.gtc_weight * gtc_loss,
+            'ctc_loss': ctc_loss,
+            'gtc_loss': gtc_loss
+        }

fastplateocr/rec/losses/ctc_loss.py

@@ -0,0 +1,33 @@
+import torch
+from torch import nn
+
+
+class CTCLoss(nn.Module):
+
+    def __init__(self, use_focal_loss=False, zero_infinity=False, **kwargs):
+        super(CTCLoss, self).__init__()
+        self.loss_func = nn.CTCLoss(blank=0,
+                                    reduction='none',
+                                    zero_infinity=zero_infinity)
+        self.use_focal_loss = use_focal_loss
+
+    def forward(self, predicts, batch):
+        # predicts = predicts['res']
+
+        batch_size = predicts.size(0)
+        label, label_length = batch[1], batch[2]
+        predicts = predicts.log_softmax(2)
+        predicts = predicts.permute(1, 0, 2)
+        preds_lengths = torch.tensor([predicts.size(0)] * batch_size,
+                                     dtype=torch.long)
+        loss = self.loss_func(predicts, label, preds_lengths, label_length)
+
+        if self.use_focal_loss:
+            # Use torch.clamp to limit the range of loss, avoiding overflow in exponential calculation
+            clamped_loss = torch.clamp(loss, min=-20, max=20)
+            weight = 1 - torch.exp(-clamped_loss)
+            weight = torch.square(weight)
+            # Use torch.where to avoid multiplying by zero weight
+            loss = torch.where(weight > 0, loss * weight, loss)
+        loss = loss.mean()
+        return {'loss': loss}

fastplateocr/rec/metrics/__init__.py

@@ -0,0 +1,33 @@
+import copy
+
+__all__ = ['build_metric']
+
+support_dict = [
+    'RecMetric', 'RecMetricLong', 'RecGTCMetric', 'RecMPGMetric', 'CMERMetric'
+]
+
+
+def build_metric(config):
+    config = copy.deepcopy(config)
+    module_name = config.pop('name')
+    assert module_name in support_dict, Exception(
+        'metric only support {}'.format(support_dict))
+
+    # Lazy import
+    if module_name == 'RecMetric':
+        from .rec_metric import RecMetric
+        module_class = RecMetric(**config)
+    elif module_name == 'RecGTCMetric':
+        from .rec_metric_gtc import RecGTCMetric
+        module_class = RecGTCMetric(**config)
+    elif module_name == 'RecMetricLong':
+        from .rec_metric_long import RecMetricLong
+        module_class = RecMetricLong(**config)
+    elif module_name == 'RecMPGMetric':
+        from .rec_metric_mgp import RecMPGMetric
+        module_class = RecMPGMetric(**config)
+    elif module_name == 'CMERMetric':
+        from .rec_metric_cmer import CMERMetric
+        module_class = CMERMetric(**config)
+
+    return module_class