litealpr 0.1.0__py3-none-any.whl

litealpr/__init__.py

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

litealpr/det/__init__.py

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

litealpr/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

litealpr/pipeline.py

@@ -0,0 +1,187 @@
+import os
+import sys
+import warnings
+
+import cv2
+import torch
+
+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 litealpr.rec.modeling import build_model
+from litealpr.rec.postprocess import build_post_process
+
+
+def download_from_hf(filename):
+    try:
+        from huggingface_hub import hf_hub_download
+        print(f"[LiteALPR] Downloading/Verifying {filename} from HuggingFace...")
+        return hf_hub_download(repo_id="anhone3/LiteALPR", filename=filename)
+    except ImportError:
+        raise ImportError("Please install huggingface_hub to auto-download models: pip install huggingface_hub")
+
+
+class LiteALPR:
+    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"[LiteALPR] 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("yolov8n_efficient/best.pt")
+
+            print(f"[LiteALPR] Loading Detection Model: {det_model_path}")
+            self.det_model = YOLO(det_model_path)
+            self.det_model.to(self.device)
+            self.det_model.fuse()
+
+        # 2. Initialize REC (SVTR26)
+        if use_rec:
+            if rec_model_path is None:
+                rec_model_path = download_from_hf("svtr26_tiny/best.pth")
+
+            print(f"[LiteALPR] 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()
+
+        if self.det_model or self.rec_model:
+            print("[LiteALPR] Models loaded successfully!")
+        else:
+            print("[LiteALPR] WARNING: No models loaded. Please provide det_model_path or rec_model_path.")
+
+    def _preprocess_crop(self, img_crop, max_ratio=12, base_shape=None, base_h=32):
+        """Preprocesses cropped image for SVTR with RatioRecTVResize logic using PIL (to match training)."""
+        from PIL import Image
+        from torchvision import transforms as T
+        from torchvision.transforms import functional as F
+
+        # Convert OpenCV BGR to PIL RGB
+        if base_shape is None:
+            base_shape = [[64, 64], [96, 48], [112, 40], [128, 32]]
+        img_rgb = cv2.cvtColor(img_crop, cv2.COLOR_BGR2RGB)
+        img = Image.fromarray(img_rgb)
+
+        w, h = img.size
+        gen_ratio = int(float(w) / float(h)) + 1
+        ratio_resize = min(gen_ratio, max_ratio)
+
+        if ratio_resize <= 4:
+            imgW, imgH = base_shape[ratio_resize - 1]
+        else:
+            imgW, imgH = [base_h * ratio_resize, base_h]
+
+        # SVTR is sensitive to interpolation algorithms; we MUST use PIL BICUBIC like during training
+        resized_image = F.resize(img, (imgH, imgW), interpolation=T.InterpolationMode.BICUBIC)
+
+        transforms = T.Compose([
+            T.ToTensor(),
+            T.Normalize(0.5, 0.5)
+        ])
+
+        tensor = transforms(resized_image)
+        tensor = tensor.unsqueeze(0)
+        return tensor.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

litealpr/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(
+            f'{module_name} is not supported. The losses in {list(name_to_module.keys())} are supportes')
+        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().__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
+        }

litealpr/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().__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}

litealpr/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(
+        f'metric only support {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

litealpr/rec/metrics/rec_metric.py

@@ -0,0 +1,287 @@
+import string
+
+import numpy as np
+from rapidfuzz.distance import Levenshtein
+
+
+def match_ss(ss1, ss2):
+    s1_len = len(ss1)
+    for c_i in range(s1_len):
+        if ss1[c_i:] == ss2[:s1_len - c_i]:
+            return ss2[s1_len - c_i:]
+    return ss2
+
+
+def stream_match(text):
+    bs = len(text)
+    s_list = []
+    conf_list = []
+    for s_conf in text:
+        s_list.append(s_conf[0])
+        conf_list.append(s_conf[1])
+    s_n = bs
+    s_start = s_list[0][:-1]
+    s_new = s_start
+    for s_i in range(1, s_n):
+        s_start = match_ss(
+            s_start, s_list[s_i][1:-1] if s_i < s_n - 1 else s_list[s_i][1:])
+        s_new += s_start
+    return s_new, sum(conf_list) / bs
+
+
+class RecMetric:
+
+    def __init__(self,
+                 main_indicator='acc',
+                 is_filter=False,
+                 is_lower=True,
+                 ignore_space=True,
+                 stream=False,
+                 with_ratio=False,
+                 max_len=25,
+                 max_ratio=4,
+                 **kwargs):
+        self.main_indicator = main_indicator
+        self.is_filter = is_filter
+        self.is_lower = is_lower
+        self.ignore_space = ignore_space
+        self.stream = stream
+        self.eps = 1e-5
+        self.with_ratio = with_ratio
+        self.max_len = max_len
+        self.max_ratio = max_ratio
+        self.reset()
+
+    def _normalize_text(self, text):
+        text = ''.join(
+            filter(lambda x: x in (string.digits + string.ascii_letters),
+                   text))
+        return text
+
+    def __call__(self,
+                 pred_label,
+                 batch=None,
+                 training=False,
+                 *args,
+                 **kwargs):
+        if self.with_ratio and not training:
+            return self.eval_all_metric(pred_label, batch)
+        else:
+            return self.eval_metric(pred_label)
+
+    def eval_metric(self, pred_label, *args, **kwargs):
+        preds, labels = pred_label
+        correct_num = 0
+        all_num = 0
+        norm_edit_dis = 0.0
+        total_edit_dis = 0.0
+        total_char_len = 0
+        for (pred, pred_conf), (target, _) in zip(preds, labels):
+            if self.stream:
+                assert len(labels) == 1
+                pred, _ = stream_match(preds)
+            if self.ignore_space:
+                pred = pred.replace(' ', '')
+                target = target.replace(' ', '')
+            if self.is_filter:
+                pred = self._normalize_text(pred)
+                target = self._normalize_text(target)
+            if self.is_lower:
+                pred = pred.lower()
+                target = target.lower()
+            
+            # For CER: absolute Levenshtein distance
+            ed = Levenshtein.distance(pred, target)
+            total_edit_dis += ed
+            total_char_len += len(target)
+            
+            norm_edit_dis += Levenshtein.normalized_distance(pred, target)
+            if pred == target:
+                correct_num += 1
+            all_num += 1
+        self.correct_num += correct_num
+        self.all_num += all_num
+        self.norm_edit_dis += norm_edit_dis
+        self.total_edit_dis += total_edit_dis
+        self.total_char_len += total_char_len
+        
+        return {
+            'acc': correct_num / (all_num + self.eps),
+            'norm_edit_dis': 1 - norm_edit_dis / (all_num + self.eps),
+            'cer': total_edit_dis / (total_char_len + self.eps),
+        }
+
+    def eval_all_metric(self, pred_label, batch=None, *args, **kwargs):
+        if self.with_ratio:
+            ratio = batch[-1]
+        preds, labels = pred_label
+        correct_num = 0
+        correct_num_real = 0
+        correct_num_lower = 0
+        correct_num_ignore_space = 0
+        correct_num_ignore_space_lower = 0
+        correct_num_ignore_space_symbol = 0
+        all_num = 0
+        norm_edit_dis = 0.0
+        total_edit_dis = 0.0
+        total_char_len = 0
+        each_len_num = [0 for _ in range(self.max_len)]
+        each_len_correct_num = [0 for _ in range(self.max_len)]
+        each_len_norm_edit_dis = [0 for _ in range(self.max_len)]
+        each_ratio_num = [0 for _ in range(self.max_ratio)]
+        each_ratio_correct_num = [0 for _ in range(self.max_ratio)]
+        each_ratio_norm_edit_dis = [0 for _ in range(self.max_ratio)]
+        for (pred, pred_conf), (target, _) in zip(preds, labels):
+            if self.stream:
+                assert len(labels) == 1
+                pred, _ = stream_match(preds)
+            if pred == target:
+                correct_num_real += 1
+
+            if pred.lower() == target.lower():
+                correct_num_lower += 1
+
+            if self.ignore_space:
+                pred = pred.replace(' ', '')
+                target = target.replace(' ', '')
+            if pred == target:
+                correct_num_ignore_space += 1
+
+            if pred.lower() == target.lower():
+                correct_num_ignore_space_lower += 1
+
+            if self.is_filter:
+                pred = self._normalize_text(pred)
+                target = self._normalize_text(target)
+            if pred == target:
+                correct_num_ignore_space_symbol += 1
+
+            if self.is_lower:
+                pred = pred.lower()
+                target = target.lower()
+                
+            ed = Levenshtein.distance(pred, target)
+            total_edit_dis += ed
+            total_char_len += len(target)
+            
+            dis = Levenshtein.normalized_distance(pred, target)
+            norm_edit_dis += dis
+            ratio_i = ratio[all_num] - 1 if ratio[
+                all_num] < self.max_ratio else self.max_ratio - 1
+            len_i = max(0, min(self.max_len, len(target)) - 1)
+            if pred == target:
+                correct_num += 1
+                each_len_correct_num[len_i] += 1
+                each_ratio_correct_num[ratio_i] += 1
+            each_len_num[len_i] += 1
+            each_len_norm_edit_dis[len_i] += dis
+
+            each_ratio_num[ratio_i] += 1
+            each_ratio_norm_edit_dis[ratio_i] += dis
+            all_num += 1
+        self.correct_num += correct_num
+        self.correct_num_real += correct_num_real
+        self.correct_num_lower += correct_num_lower
+        self.correct_num_ignore_space += correct_num_ignore_space
+        self.correct_num_ignore_space_lower += correct_num_ignore_space_lower
+        self.correct_num_ignore_space_symbol += correct_num_ignore_space_symbol
+        self.all_num += all_num
+        self.norm_edit_dis += norm_edit_dis
+        self.total_edit_dis += total_edit_dis
+        self.total_char_len += total_char_len
+        self.each_len_num = self.each_len_num + np.array(each_len_num)
+        self.each_len_correct_num = self.each_len_correct_num + np.array(
+            each_len_correct_num)
+        self.each_len_norm_edit_dis = self.each_len_norm_edit_dis + np.array(
+            each_len_norm_edit_dis)
+        self.each_ratio_num = self.each_ratio_num + np.array(each_ratio_num)
+        self.each_ratio_correct_num = self.each_ratio_correct_num + np.array(
+            each_ratio_correct_num)
+        self.each_ratio_norm_edit_dis = self.each_ratio_norm_edit_dis + np.array(
+            each_ratio_norm_edit_dis)
+        return {
+            'acc': correct_num / (all_num + self.eps),
+            'norm_edit_dis': 1 - norm_edit_dis / (all_num + self.eps),
+            'cer': total_edit_dis / (total_char_len + self.eps),
+        }
+
+    def get_metric(self, training=False):
+        if self.with_ratio and not training:
+            return self.get_all_metric()
+        acc = 1.0 * self.correct_num / (self.all_num + self.eps)
+        norm_edit_dis = 1 - self.norm_edit_dis / (self.all_num + self.eps)
+        cer = self.total_edit_dis / (self.total_char_len + self.eps)
+        num_samples = self.all_num
+        self.reset()
+        return {
+            'acc': acc,
+            'norm_edit_dis': norm_edit_dis,
+            'cer': cer,
+            'num_samples': num_samples
+        }
+
+    def get_all_metric(self):
+        acc = 1.0 * self.correct_num / (self.all_num + self.eps)
+        acc_real = 1.0 * self.correct_num_real / (self.all_num + self.eps)
+        acc_lower = 1.0 * self.correct_num_lower / (self.all_num + self.eps)
+        acc_ignore_space = 1.0 * self.correct_num_ignore_space / (
+            self.all_num + self.eps)
+        acc_ignore_space_lower = 1.0 * self.correct_num_ignore_space_lower / (
+            self.all_num + self.eps)
+        acc_ignore_space_symbol = 1.0 * self.correct_num_ignore_space_symbol / (
+            self.all_num + self.eps)
+
+        norm_edit_dis = 1 - self.norm_edit_dis / (self.all_num + self.eps)
+        cer = self.total_edit_dis / (self.total_char_len + self.eps)
+        num_samples = self.all_num
+        each_len_acc = (self.each_len_correct_num /
+                        (self.each_len_num + self.eps)).tolist()
+        each_len_norm_edit_dis = (1 -
+                                  ((self.each_len_norm_edit_dis) /
+                                   ((self.each_len_num) + self.eps))).tolist()
+        each_len_num = self.each_len_num.tolist()
+        each_ratio_acc = (self.each_ratio_correct_num /
+                          (self.each_ratio_num + self.eps)).tolist()
+        each_ratio_norm_edit_dis = (1 - ((self.each_ratio_norm_edit_dis) / (
+            (self.each_ratio_num) + self.eps))).tolist()
+        each_ratio_num = self.each_ratio_num.tolist()
+        self.reset()
+        return {
+            'acc': acc,
+            'acc_real': acc_real,
+            'acc_lower': acc_lower,
+            'acc_ignore_space': acc_ignore_space,
+            'acc_ignore_space_lower': acc_ignore_space_lower,
+            'acc_ignore_space_symbol': acc_ignore_space_symbol,
+            'acc_ignore_space_lower_symbol': acc,
+            'each_len_num': each_len_num,
+            'each_len_acc': each_len_acc,
+            'each_len_norm_edit_dis': each_len_norm_edit_dis,
+            'each_ratio_num': each_ratio_num,
+            'each_ratio_acc': each_ratio_acc,
+            'each_ratio_norm_edit_dis': each_ratio_norm_edit_dis,
+            'norm_edit_dis': norm_edit_dis,
+            'cer': cer,
+            'num_samples': num_samples
+        }
+
+    def reset(self):
+        self.correct_num = 0
+        self.all_num = 0
+        self.norm_edit_dis = 0
+        self.total_edit_dis = 0.0
+        self.total_char_len = 0
+        self.correct_num_real = 0
+        self.correct_num_lower = 0
+        self.correct_num_ignore_space = 0
+        self.correct_num_ignore_space_lower = 0
+        self.correct_num_ignore_space_symbol = 0
+        self.each_len_num = np.array([0 for _ in range(self.max_len)])
+        self.each_len_correct_num = np.array([0 for _ in range(self.max_len)])
+        self.each_len_norm_edit_dis = np.array(
+            [0. for _ in range(self.max_len)])
+        self.each_ratio_num = np.array([0 for _ in range(self.max_ratio)])
+        self.each_ratio_correct_num = np.array(
+            [0 for _ in range(self.max_ratio)])
+        self.each_ratio_norm_edit_dis = np.array(
+            [0. for _ in range(self.max_ratio)])

litealpr/rec/modeling/__init__.py

@@ -0,0 +1,11 @@
+import copy
+
+from .base_recognizer import BaseRecognizer
+
+__all__ = ['build_model']
+
+
+def build_model(config):
+    config = copy.deepcopy(config)
+    rec_model = BaseRecognizer(config)
+    return rec_model