asterocr 1.0.0__py3-none-any.whl

asterocr/__init__.py

@@ -0,0 +1,12 @@
+"""
+AsterOCR - Modern High-Performance Vietnamese OCR Engine
+Designed to replace VietOCR with non-autoregressive CTC sequence modeling,
+dynamic synthetic data generation, and sub-second recognition.
+"""
+
+__version__ = "1.0.0"
+
+from asterocr.pipeline import AsterOCR
+from asterocr.vocab import default_vocab
+
+__all__ = ["AsterOCR", "default_vocab", "__version__"]

asterocr/cli.py

@@ -0,0 +1,35 @@
+"""
+Command Line Interface for AsterOCR
+Usage:
+    asterocr path/to/image.jpg
+    asterocr path/to/image.jpg --backend onnx
+"""
+
+import sys
+import argparse
+import os
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="asterocr",
+        description="AsterOCR: Ultra-fast Vietnamese Text Recognition CLI"
+    )
+    parser.add_argument("image", help="Path to input image file (jpg, png, webp)")
+    parser.add_argument("--backend", default="onnx", choices=["onnx", "torch"], help="Inference backend (default: onnx)")
+    parser.add_argument("--device", default="cpu", help="Computation device (cpu or cuda)")
+    parser.add_argument("--model-path", default=None, help="Explicit path to asternet.onnx or asternet.pth")
+    parser.add_argument("--version", action="version", version="AsterOCR v1.0.0")
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.image):
+        print(f"[-] Lỗi: File ảnh không tồn tại: {args.image}", file=sys.stderr)
+        sys.exit(1)
+
+    from asterocr.pipeline import AsterOCR
+    ocr = AsterOCR(model_path=args.model_path, backend=args.backend, device=args.device)
+    result = ocr.predict(args.image)
+    print(result)
+
+if __name__ == "__main__":
+    main()

asterocr/comparator.py

@@ -0,0 +1,412 @@
+"""
+AsterOCR vs VietOCR Comparison Engine
+Provides simultaneous side-by-side inference, precision comparison,
+latency benchmark, and full page document parsing.
+"""
+
+import os
+import sys
+import time
+import io
+import json
+import base64
+import torch
+import cv2
+import numpy as np
+from PIL import Image
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from asterocr.inference.predictor import AsterPredictor
+from asterocr.trainer.metrics import calculate_cer, calculate_wer, levenshtein_distance
+from vietocr.tool.config import Cfg
+from vietocr.tool.predictor import Predictor
+from rapidocr_onnxruntime import RapidOCR
+
+class OCRComparator:
+    _instance = None
+
+    @classmethod
+    def get_instance(cls):
+        if cls._instance is None:
+            cls._instance = OCRComparator()
+        return cls._instance
+
+    def __init__(self):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"[*] OCRComparator khởi tạo trên thiết bị: {self.device}")
+        
+        # 1. AsterNet SOTA Predictor
+        model_path = os.path.join(os.path.dirname(__file__), "..", "models", "asternet_best.pth")
+        if not os.path.exists(model_path):
+            model_path = os.path.join("models", "asternet_best.pth")
+        self.model_path = os.path.abspath(model_path)
+        self.model_mtime = os.path.getmtime(self.model_path) if os.path.exists(self.model_path) else 0
+        self.aster_pred = AsterPredictor(model_path=self.model_path, device=self.device)
+        print(f"[+] AsterNet nạp thành công từ: {self.model_path}")
+        
+        # 2. VietOCR Predictor
+        config = Cfg.load_config_from_name('vgg_transformer')
+        config['device'] = self.device
+        config['predictor']['beamsearch'] = False
+        self.viet_pred = Predictor(config)
+        print("[+] VietOCR (vgg_transformer) nạp thành công!")
+
+        # 3. Detector cho Full-page document
+        try:
+            self.detector = RapidOCR(use_cls=False, use_rec=False)
+            print("[+] RapidOCR text line detector sẵn sàng!")
+        except Exception as e:
+            print(f"[!] Cảnh báo khởi tạo RapidOCR: {e}")
+            self.detector = None
+
+        # Benchmark running state
+        self.benchmark_state = {
+            "status": "idle",
+            "progress": 0,
+            "total": 0,
+            "current": 0,
+            "dataset_name": "",
+            "metrics": {},
+            "results": [],
+            "error": None
+        }
+
+    def check_and_reload(self):
+        """Hot-reloads AsterNet if weights file has been updated by training."""
+        if os.path.exists(self.model_path):
+            cur_mtime = os.path.getmtime(self.model_path)
+            if cur_mtime > self.model_mtime:
+                try:
+                    print(f"[*] Phát hiện weights mới ({time.ctime(cur_mtime)}), đang tự động nạp lại AsterNet...")
+                    from asterocr.inference.predictor import AsterPredictor
+                    self.aster_pred = AsterPredictor(model_path=self.model_path, device=self.device)
+                    self.model_mtime = cur_mtime
+                    print("[+] Tự động nạp lại AsterNet thành công!")
+                except Exception as e:
+                    print(f"[!] Lỗi khi nạp lại AsterNet: {e}")
+
+    def compare_crop(self, pil_img, gt_text=None):
+        """
+        Runs both models on a single line crop and returns side-by-side metrics.
+        """
+        self.check_and_reload()
+        # Warmup/ensure RGB
+        if pil_img.mode != "RGB":
+            pil_img = pil_img.convert("RGB")
+
+        # AsterNet inference
+        t0 = time.perf_counter()
+        aster_text = self.aster_pred.predict(pil_img)
+        aster_lat = (time.perf_counter() - t0) * 1000
+
+        # VietOCR inference
+        t0 = time.perf_counter()
+        viet_text = self.viet_pred.predict(pil_img)
+        viet_lat = (time.perf_counter() - t0) * 1000
+
+        speedup = viet_lat / max(0.001, aster_lat)
+
+        res = {
+            "aster": {
+                "text": aster_text,
+                "latency_ms": round(aster_lat, 2),
+                "fps": round(1000.0 / max(1.0, aster_lat), 1)
+            },
+            "vietocr": {
+                "text": viet_text,
+                "latency_ms": round(viet_lat, 2),
+                "fps": round(1000.0 / max(1.0, viet_lat), 1)
+            },
+            "speedup": round(speedup, 2)
+        }
+
+        if gt_text is not None and len(gt_text.strip()) > 0:
+            gt = gt_text.strip()
+            aster_cer = calculate_cer([gt], [aster_text]) * 100
+            aster_wer = calculate_wer([gt], [aster_text]) * 100
+            viet_cer = calculate_cer([gt], [viet_text]) * 100
+            viet_wer = calculate_wer([gt], [viet_text]) * 100
+            
+            res["gt"] = gt
+            res["aster"]["cer"] = round(aster_cer, 2)
+            res["aster"]["wer"] = round(aster_wer, 2)
+            res["aster"]["match"] = (aster_text.strip() == gt)
+            
+            res["vietocr"]["cer"] = round(viet_cer, 2)
+            res["vietocr"]["wer"] = round(viet_wer, 2)
+            res["vietocr"]["match"] = (viet_text.strip() == gt)
+
+            if aster_cer < viet_cer:
+                winner = "asternet"
+            elif viet_cer < aster_cer:
+                winner = "vietocr"
+            else:
+                winner = "draw"
+            res["accuracy_winner"] = winner
+
+        return res
+
+    def compare_document_page(self, np_img_bgr):
+        """
+        Detects all text lines on a document page and recognizes with both models.
+        Returns visual annotated image base64 and line-by-line comparison data.
+        """
+        self.check_and_reload()
+        if self.detector is None:
+            return {"error": "RapidOCR detector không khả dụng."}
+
+        boxes, _ = self.detector(np_img_bgr)
+        if not boxes:
+            return {"lines": [], "annotated_image": None, "total_lines": 0}
+
+        h_img, w_img = np_img_bgr.shape[:2]
+        lines_info = []
+        pil_crops = []
+        box_coords = []
+
+        # Sort boxes top to bottom
+        sorted_boxes = sorted(boxes, key=lambda b: np.min(np.array(b[0] if isinstance(b[0], list) and len(b) == 2 else b)[:, 1]))
+
+        for idx, b in enumerate(sorted_boxes):
+            pts = np.array(b[0] if isinstance(b[0], list) and len(b) == 2 else b)
+            x_min, y_min = np.min(pts, axis=0)
+            x_max, y_max = np.max(pts, axis=0)
+
+            pad = 2
+            x1 = max(0, int(x_min) - pad)
+            y1 = max(0, int(y_min) - pad)
+            x2 = min(w_img, int(x_max) + pad)
+            y2 = min(h_img, int(y_max) + pad)
+
+            if y2 - y1 < 6 or x2 - x1 < 6:
+                continue
+
+            crop = np_img_bgr[y1:y2, x1:x2]
+            if crop.size == 0:
+                continue
+
+            crop_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
+            pil_crops.append(crop_pil)
+            box_coords.append((x1, y1, x2, y2))
+
+        # Annotate image
+        annotated = np_img_bgr.copy()
+        
+        # Batch inference
+        t0 = time.perf_counter()
+        aster_preds = [self.aster_pred.predict(c) for c in pil_crops]
+        aster_total_time = (time.perf_counter() - t0) * 1000
+
+        t0 = time.perf_counter()
+        viet_preds = [self.viet_pred.predict(c) for c in pil_crops]
+        viet_total_time = (time.perf_counter() - t0) * 1000
+
+        for i, (coords, ap, vp) in enumerate(zip(box_coords, aster_preds, viet_preds)):
+            x1, y1, x2, y2 = coords
+            # Draw bbox: Cyan for line
+            cv2.rectangle(annotated, (x1, y1), (x2, y2), (0, 255, 255), 1)
+            cv2.putText(annotated, str(i + 1), (x1, max(12, y1 - 2)), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 255, 255), 1)
+
+            # Crop thumbnail base64
+            crop_rgb = pil_crops[i]
+            buf = io.BytesIO()
+            crop_rgb.save(buf, format="JPEG", quality=85)
+            crop_b64 = "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
+
+            lines_info.append({
+                "line_id": i + 1,
+                "bbox": [x1, y1, x2, y2],
+                "crop_thumbnail": crop_b64,
+                "aster_text": ap,
+                "viet_text": vp,
+                "is_identical": (ap.strip() == vp.strip())
+            })
+
+        # Encode annotated image
+        _, enc_jpg = cv2.imencode(".jpg", annotated, [cv2.IMWRITE_JPEG_QUALITY, 80])
+        annotated_b64 = "data:image/jpeg;base64," + base64.b64encode(enc_jpg).decode()
+
+        return {
+            "total_lines": len(lines_info),
+            "aster_total_ms": round(aster_total_time, 2),
+            "viet_total_ms": round(viet_total_time, 2),
+            "speedup": round(viet_total_time / max(0.001, aster_total_time), 2),
+            "lines": lines_info,
+            "annotated_image": annotated_b64
+        }
+
+    def start_batch_benchmark(self, dataset_name, max_samples=100):
+        """
+        Runs benchmark in a background thread and updates self.benchmark_state.
+        """
+        import threading
+        if self.benchmark_state["status"] == "running":
+            return False, "Benchmark đang chạy. Hãy đợi hoặc dừng phiên hiện tại."
+
+        self.benchmark_state = {
+            "status": "running",
+            "progress": 0,
+            "total": 0,
+            "current": 0,
+            "dataset_name": dataset_name,
+            "metrics": {},
+            "results": [],
+            "error": None
+        }
+
+        t = threading.Thread(target=self._run_benchmark_worker, args=(dataset_name, max_samples), daemon=True)
+        t.start()
+        return True, "Benchmark đã bắt đầu."
+
+    def _run_benchmark_worker(self, dataset_name, max_samples):
+        try:
+            # Hot-reload linguistics and weights before running benchmark
+            try:
+                import importlib
+                import asterocr.inference.viet_linguistics as vl
+                importlib.reload(vl)
+                self.check_and_reload()
+            except Exception as e:
+                print(f"[!] Warning reloading modules: {e}")
+
+            samples = []
+            if dataset_name == "mega_test":
+                labels_file = os.path.join("dataset_mega_test", "labels.json")
+                if os.path.exists(labels_file):
+                    with open(labels_file, "r", encoding="utf-8") as f:
+                        meta = json.load(f)
+                    for item in meta:
+                        img_p = os.path.join("dataset_mega_test", item["image_path"])
+                        if os.path.exists(img_p):
+                            img = Image.open(img_p).convert("RGB")
+                            samples.append((img, item["text"], os.path.basename(img_p)))
+            elif dataset_name in ["real_scans", "unified", "viet_real"]:
+                if dataset_name == "viet_real":
+                    lmdb_path = "dataset_viet_real.lmdb"
+                elif dataset_name == "unified":
+                    lmdb_path = "dataset_unified.lmdb"
+                else:
+                    lmdb_path = "dataset_real_scans.lmdb"
+                if os.path.exists(lmdb_path):
+                    import lmdb
+                    env = lmdb.open(lmdb_path, readonly=True, max_readers=1)
+                    with env.begin() as txn:
+                        num_val = txn.get(b'num-samples')
+                        total_db = int(num_val.decode('utf-8')) if num_val else 0
+                        step = max(1, total_db // max(1, max_samples))
+                        for i in range(1, total_db, step):
+                            if len(samples) >= max_samples:
+                                break
+                            img_key = f"image-{i:09d}".encode()
+                            lbl_key = f"label-{i:09d}".encode()
+                            img_b = txn.get(img_key)
+                            lbl_b = txn.get(lbl_key)
+                            if img_b and lbl_b:
+                                try:
+                                    img = Image.open(io.BytesIO(img_b)).convert("RGB")
+                                    txt = lbl_b.decode('utf-8').strip()
+                                    if txt:
+                                        samples.append((img, txt, f"sample_{i:07d}.jpg"))
+                                except Exception:
+                                    pass
+                    env.close()
+
+            if not samples:
+                self.benchmark_state["status"] = "error"
+                self.benchmark_state["error"] = f"Không tìm thấy dữ liệu cho '{dataset_name}'."
+                return
+
+            self.benchmark_state["total"] = len(samples)
+            results = []
+            gt_all = []
+            aster_all = []
+            viet_all = []
+
+            t_aster_total = 0.0
+            t_viet_total = 0.0
+
+            for idx, (img, gt, name) in enumerate(samples):
+                # AsterNet
+                t0 = time.perf_counter()
+                ap = self.aster_pred.predict(img)
+                t_a = (time.perf_counter() - t0) * 1000
+                t_aster_total += t_a
+
+                # VietOCR
+                t0 = time.perf_counter()
+                vp = self.viet_pred.predict(img)
+                t_v = (time.perf_counter() - t0) * 1000
+                t_viet_total += t_v
+
+                gt_all.append(gt)
+                aster_all.append(ap)
+                viet_all.append(vp)
+
+                # Thumbnail
+                buf = io.BytesIO()
+                img.save(buf, format="JPEG", quality=80)
+                img_b64 = "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()
+
+                # Per-sample metrics
+                a_match = (ap.strip() == gt.strip())
+                v_match = (vp.strip() == gt.strip())
+
+                results.append({
+                    "id": idx + 1,
+                    "name": name,
+                    "thumbnail": img_b64,
+                    "gt": gt,
+                    "aster_pred": ap,
+                    "viet_pred": vp,
+                    "aster_match": a_match,
+                    "viet_match": v_match,
+                    "aster_ms": round(t_a, 1),
+                    "viet_ms": round(t_v, 1),
+                    "aster_cer": round(calculate_cer([gt], [ap]) * 100, 1),
+                    "viet_cer": round(calculate_cer([gt], [vp]) * 100, 1)
+                })
+
+                self.benchmark_state["current"] = idx + 1
+                self.benchmark_state["progress"] = round((idx + 1) / len(samples) * 100, 1)
+
+            # Final metrics
+            N = len(samples)
+            aster_cer = round(calculate_cer(gt_all, aster_all) * 100, 2)
+            aster_wer = round(calculate_wer(gt_all, aster_all) * 100, 2)
+            aster_exact = round(sum(1 for g, p in zip(gt_all, aster_all) if g.strip() == p.strip()) / N * 100, 2)
+            aster_fps = round(N / (t_aster_total / 1000.0), 1)
+            aster_avg_ms = round(t_aster_total / N, 2)
+
+            viet_cer = round(calculate_cer(gt_all, viet_all) * 100, 2)
+            viet_wer = round(calculate_wer(gt_all, viet_all) * 100, 2)
+            viet_exact = round(sum(1 for g, p in zip(gt_all, viet_all) if g.strip() == p.strip()) / N * 100, 2)
+            viet_fps = round(N / (t_viet_total / 1000.0), 1)
+            viet_avg_ms = round(t_viet_total / N, 2)
+
+            speedup = round(aster_fps / max(0.01, viet_fps), 2)
+
+            self.benchmark_state["metrics"] = {
+                "aster": {
+                    "cer": aster_cer,
+                    "wer": aster_wer,
+                    "exact": aster_exact,
+                    "fps": aster_fps,
+                    "avg_ms": aster_avg_ms
+                },
+                "vietocr": {
+                    "cer": viet_cer,
+                    "wer": viet_wer,
+                    "exact": viet_exact,
+                    "fps": viet_fps,
+                    "avg_ms": viet_avg_ms
+                },
+                "speedup": speedup,
+                "winner": "AsterNet" if aster_cer <= viet_cer else "VietOCR"
+            }
+            self.benchmark_state["results"] = results
+            self.benchmark_state["status"] = "finished"
+
+        except Exception as e:
+            self.benchmark_state["status"] = "error"
+            self.benchmark_state["error"] = str(e)

asterocr/dataset/export_real_zip.py

@@ -0,0 +1,117 @@
+"""
+AsterOCR Fast Real Dataset Exporter
+Extracts unexported real document scan samples (en_00, en_01, random) from ocr_dataset.zip
+and packages them into dataset_real_scans.lmdb for seamless streaming to Dual GPUs.
+"""
+
+import os
+import io
+import sys
+import json
+import time
+import zipfile
+import lmdb
+from PIL import Image
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
+def export_real_zip_to_lmdb(
+    zip_path="dataset_download/ocr_dataset.zip",
+    index_path="dataset_download/ocr_dataset_index.json",
+    output_dir="dataset_real_scans.lmdb",
+    target_height=32,
+    categories=("en_00", "en_01", "random"),
+    batch_commit=5000
+):
+    if not os.path.exists(zip_path):
+        print(f"[!] Không tìm thấy tệp zip: {zip_path}")
+        return False
+    if not os.path.exists(index_path):
+        print(f"[!] Không tìm thấy tệp index: {index_path}")
+        return False
+
+    print(f"[*] Đang tải danh mục từ {index_path}...")
+    with open(index_path, "r", encoding="utf-8") as f:
+        all_entries = json.load(f)
+
+    # Filter targets
+    target_entries = [
+        item for item in all_entries
+        if any(item[0].startswith(cat + "/") for cat in categories)
+    ]
+    total_target = len(target_entries)
+    print(f"[+] Tìm thấy {total_target:,} mẫu ảnh thật thuộc các nhóm: {categories}")
+
+    os.makedirs(output_dir, exist_ok=True)
+    # 20GB map size ceiling
+    env = lmdb.open(output_dir, map_size=25 * 1024 * 1024 * 1024, sync=False)
+
+    t0 = time.time()
+    written_count = 0
+    txn = env.begin(write=True)
+
+    with zipfile.ZipFile(zip_path, "r") as zf:
+        for idx, (img_rel, txt_rel) in enumerate(target_entries, 1):
+            try:
+                img_data = zf.read(img_rel)
+                txt_data = zf.read(txt_rel).decode("utf-8", errors="ignore").strip()
+                if not txt_data:
+                    continue
+
+                # Resize to target height 32
+                im = Image.open(io.BytesIO(img_data)).convert("RGB")
+                w, h = im.size
+                if w < 4 or h < 4:
+                    continue
+                new_w = max(16, int((float(target_height) / max(1, h)) * w))
+                im_res = im.resize((new_w, target_height), Image.BILINEAR)
+
+                buf = io.BytesIO()
+                im_res.save(buf, format="JPEG", quality=90)
+                jpeg_bytes = buf.getvalue()
+
+                written_count += 1
+                img_key = f"image-{written_count:09d}".encode("utf-8")
+                lbl_key = f"label-{written_count:09d}".encode("utf-8")
+
+                txn.put(img_key, jpeg_bytes)
+                txn.put(lbl_key, txt_data.encode("utf-8"))
+
+                if written_count % batch_commit == 0:
+                    txn.commit()
+                    txn = env.begin(write=True)
+                    elapsed = time.time() - t0
+                    sps = written_count / max(0.1, elapsed)
+                    print(f"--> Đã xuất: {written_count:,}/{total_target:,} mẫu ({written_count/total_target*100:.1f}%) | {sps:.0f} mẫu/s")
+
+            except Exception as e:
+                continue
+
+    # Finalize
+    txn.put(b"num-samples", str(written_count).encode("utf-8"))
+    txn.commit()
+    env.sync()
+    env.close()
+
+    total_time = time.time() - t0
+    info = {
+        "format": "lmdb",
+        "num_samples": written_count,
+        "categories": list(categories),
+        "source": zip_path,
+        "created_at": time.strftime("%Y-%m-%d %H:%M:%S"),
+        "total_seconds": round(total_time, 1)
+    }
+    with open(os.path.join(output_dir, "dataset_info.json"), "w", encoding="utf-8") as f:
+        json.dump(info, f, indent=2, ensure_ascii=False)
+
+    print(f"\n[+] HOÀN TẤT ĐÓNG GÓI {written_count:,} MẪU VÀO {output_dir} TRONG {total_time:.1f} GIÂY!")
+    return True
+
+if __name__ == "__main__":
+    export_real_zip_to_lmdb()