asterocr 1.0.0__py3-none-any.whl

asterocr/dataset/unified_packer.py

@@ -0,0 +1,402 @@
+"""
+AsterOCR Multi-Source Unified Dataset Packer (Industrial LMDB)
+Merges multiple diverse real and synthetic data streams into a single unified LMDB database:
+1. Real Vietnamese Handwriting from VNOnDB (InkData_line_processed in ocr_dataset.zip)
+2. Real Receipts, Bills, Invoices & Vouchers (meta in ocr_dataset.zip)
+3. Real Scanned Documents (vi_00, vi_01 in ocr_dataset.zip)
+4. Millions of Synthetic Realistic Vietnamese Handwriting Lines (notebook paper, pen pressure, slants)
+5. Millions of Synthetic Realistic Thermal Receipts (dot-matrix, store names, VAT, amounts)
+6. Open-Source Hugging Face Vietnamese Handwriting datasets
+
+Output Format:
+- Single LMDB directory: dataset_unified.lmdb (data.mdb)
+- image-%09d: compressed JPEG bytes (zero NTFS file fragmentation)
+- label-%09d: UTF-8 Vietnamese ground-truth string
+- source-%09d: source tag string (real_handwriting, real_receipt, synth_handwriting, etc.)
+- num-samples: total number of samples
+- dataset_summary.json: complete category and distribution report
+"""
+
+import os
+import io
+import sys
+import json
+import time
+import random
+import zipfile
+import argparse
+
+# Suppress OpenBLAS / MKL multi-threading collisions in child worker processes
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+
+from multiprocessing import Pool, cpu_count
+import numpy as np
+from PIL import Image
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
+try:
+    import lmdb
+except ImportError:
+    lmdb = None
+
+# Ensure root in sys.path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))
+from asterocr.synthesizer.handwriting_generator import (
+    RealisticHandwritingRenderer,
+    RealisticReceiptRenderer,
+    sample_handwritten_text,
+    sample_receipt_text,
+    _worker_chunk_handwriting
+)
+
+def _read_real_zip_sample(zip_ref, img_path, txt_path, target_height=32):
+    """Safely extracts, resizes, and encodes a sample from the real OCR zip."""
+    try:
+        img_bytes = zip_ref.read(img_path)
+        txt_bytes = zip_ref.read(txt_path)
+        text = txt_bytes.decode('utf-8', errors='ignore').strip()
+        if not text:
+            return None, None
+            
+        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+        w, h = img.size
+        if w < 4 or h < 4:
+            return None, None
+            
+        new_w = max(16, int((float(target_height) / max(1, h)) * w))
+        resized_img = img.resize((new_w, target_height), Image.LANCZOS)
+        
+        buf = io.BytesIO()
+        resized_img.save(buf, format="JPEG", quality=92)
+        return buf.getvalue(), text
+    except Exception:
+        return None, None
+
+def pack_unified_dataset(
+    output_dir="dataset_unified.lmdb",
+    zip_path="dataset_download/ocr_dataset.zip",
+    index_path="dataset_download/ocr_dataset_index.json",
+    num_synth_handwriting=100000,
+    num_synth_receipt=50000,
+    include_real_zip=True,
+    max_real_handwriting=None,
+    max_real_receipts=None,
+    max_real_documents=None,
+    hf_parquet_files=None,
+    num_workers=10,
+    target_height=32
+):
+    """
+    Master pipeline to assemble the unified OCR dataset.
+    """
+    if lmdb is None:
+        raise ImportError("Gói 'lmdb' chưa được cài đặt! Hãy chạy: pip install lmdb")
+
+    num_workers = min(10, max(1, num_workers)) # Cap at 10 workers for CPU stability
+    os.makedirs(output_dir, exist_ok=True)
+    
+    print("=" * 72)
+    print("   ASTER-OCR UNIFIED MULTI-SOURCE DATASET PACKER (LMDB)")
+    print("=" * 72)
+    print(f"  Thư mục đích           : {output_dir}")
+    print(f"  Số luồng CPU tối đa    : {num_workers} workers")
+    print(f"  Chiều cao ảnh chuẩn hóa: {target_height} px")
+    print(f"  Chữ viết tay giả lập   : {num_synth_handwriting:,} mẫu")
+    print(f"  Biên lai giả lập       : {num_synth_receipt:,} mẫu")
+    print("=" * 72)
+
+    # 1. Estimate total samples to size the LMDB map
+    total_estimated = num_synth_handwriting + num_synth_receipt
+    real_items = []
+    
+    if include_real_zip and os.path.exists(zip_path) and os.path.exists(index_path):
+        print(f"[*] Đang tải chỉ mục dữ liệu thực tế từ: {index_path}...")
+        try:
+            with open(index_path, "r", encoding="utf-8") as f:
+                all_index = json.load(f)
+            
+            # Categorize real items
+            hw_items = [x for x in all_index if "InkData" in x[0]]
+            rc_items = [x for x in all_index if x[0].startswith("meta/")]
+            doc_items = [x for x in all_index if x[0].startswith(("vi_00/", "vi_01/"))]
+            other_items = [x for x in all_index if x[0].startswith(("random/", "en_00/", "en_01/"))]
+            
+            print(f"    - Viết tay thật VNOnDB : {len(hw_items):,} mẫu")
+            print(f"    - Hóa đơn / biên lai   : {len(rc_items):,} mẫu")
+            print(f"    - Tài liệu in tiếng Việt: {len(doc_items):,} mẫu")
+            print(f"    - Mã số & song ngữ     : {len(other_items):,} mẫu")
+            
+            # Slice according to quotas
+            if max_real_handwriting is not None:
+                hw_items = hw_items[:max_real_handwriting]
+            if max_real_receipts is not None:
+                rc_items = rc_items[:max_real_receipts]
+            if max_real_documents is not None:
+                doc_items = doc_items[:max_real_documents]
+                
+            real_items.extend([(item, "real_handwriting") for item in hw_items])
+            real_items.extend([(item, "real_receipt") for item in rc_items])
+            real_items.extend([(item, "real_document") for item in doc_items])
+            
+            total_estimated += len(real_items)
+            print(f"[+] Sẽ gom {len(real_items):,} mẫu thực tế từ kho zip!")
+        except Exception as e:
+            print(f"[-] Lỗi khi nạp chỉ mục zip: {e}")
+
+    # Map size: ~35KB per sample, minimum 64MB, with 1.3x buffer
+    map_size = max(1024 * 1024 * 64, int(total_estimated * 45 * 1024))
+    
+    print(f"[*] Khởi tạo LMDB môi trường: map_size = {map_size // (1024 * 1024):,} MB...")
+    env = lmdb.open(output_dir, map_size=map_size)
+    txn = env.begin(write=True)
+    
+    samples_written = 0
+    uncommitted = 0
+    commit_interval = 5000
+    start_time = time.time()
+    
+    stats = {
+        "real_handwriting": 0,
+        "real_receipt": 0,
+        "real_document": 0,
+        "synth_handwriting": 0,
+        "synth_receipt": 0,
+        "hf_handwriting": 0,
+    }
+
+    try:
+        # -------------------------------------------------------------
+        # PHASE 1: Stream Real Images from ZIP directly into LMDB
+        # -------------------------------------------------------------
+        if real_items and os.path.exists(zip_path):
+            print("\n[GIAI ĐOẠN 1] Bắt đầu nạp dữ liệu thực tế (Viết tay thật + Hóa đơn + Tài liệu)...")
+            with zipfile.ZipFile(zip_path, "r") as z:
+                for idx, (item, source_tag) in enumerate(real_items):
+                    img_path, txt_path = item
+                    img_bytes, text = _read_real_zip_sample(z, img_path, txt_path, target_height)
+                    if img_bytes is None or not text:
+                        continue
+                        
+                    samples_written += 1
+                    img_key = f"image-{samples_written:09d}".encode("utf-8")
+                    lbl_key = f"label-{samples_written:09d}".encode("utf-8")
+                    src_key = f"source-{samples_written:09d}".encode("utf-8")
+                    
+                    txn.put(img_key, img_bytes)
+                    txn.put(lbl_key, text.encode("utf-8"))
+                    txn.put(src_key, source_tag.encode("utf-8"))
+                    
+                    stats[source_tag] = stats.get(source_tag, 0) + 1
+                    uncommitted += 1
+                    
+                    if uncommitted >= commit_interval:
+                        txn.commit()
+                        txn = env.begin(write=True)
+                        uncommitted = 0
+                        
+                    if samples_written % 1000 == 0 or idx == len(real_items) - 1:
+                        elapsed = max(0.001, time.time() - start_time)
+                        speed = samples_written / elapsed
+                        sys.stdout.write(
+                            f"\r[*] [Real Data] Đã nạp: {samples_written:,} mẫu | "
+                            f"Tốc độ: {speed:,.0f} mẫu/s  "
+                        )
+                        sys.stdout.flush()
+                        
+            print(f"\n[+] Đã nạp xong toàn bộ dữ liệu thực tế!")
+
+        # -------------------------------------------------------------
+        # PHASE 2: Parallel Synthetic Generation (Handwriting & Receipts)
+        # -------------------------------------------------------------
+        synth_tasks = []
+        chunk_size = 1000
+        
+        # Handwriting tasks
+        curr_cnt = 0
+        while curr_cnt < num_synth_handwriting:
+            cnt = min(chunk_size, num_synth_handwriting - curr_cnt)
+            synth_tasks.append((samples_written + 1, cnt, "handwriting", target_height))
+            curr_cnt += cnt
+            samples_written += cnt  # pre-allocate ID space
+            
+        # Receipt tasks
+        curr_rc_cnt = 0
+        while curr_rc_cnt < num_synth_receipt:
+            cnt = min(chunk_size, num_synth_receipt - curr_rc_cnt)
+            synth_tasks.append((samples_written + 1, cnt, "receipt", target_height))
+            curr_rc_cnt += cnt
+            samples_written += cnt
+            
+        # Reset actual counter for writing
+        samples_written -= (num_synth_handwriting + num_synth_receipt)
+        
+        if synth_tasks:
+            print(f"\n[GIAI ĐOẠN 2] Sinh dữ liệu viết tay & hóa đơn song song ({num_workers} workers)...")
+            synth_start_time = time.time()
+            with Pool(num_workers) as pool:
+                for chunk_data in pool.imap_unordered(_worker_chunk_handwriting, synth_tasks, chunksize=1):
+                    for _, img_bytes, text, source_type in chunk_data:
+                        samples_written += 1
+                        img_key = f"image-{samples_written:09d}".encode("utf-8")
+                        lbl_key = f"label-{samples_written:09d}".encode("utf-8")
+                        src_key = f"source-{samples_written:09d}".encode("utf-8")
+                        
+                        txn.put(img_key, img_bytes)
+                        txn.put(lbl_key, text.encode("utf-8"))
+                        txn.put(src_key, source_type.encode("utf-8"))
+                        
+                        stats[source_type] = stats.get(source_type, 0) + 1
+                        uncommitted += 1
+                        
+                        if uncommitted >= commit_interval:
+                            txn.commit()
+                            txn = env.begin(write=True)
+                            uncommitted = 0
+                            
+                    elapsed = max(0.001, time.time() - start_time)
+                    speed = samples_written / elapsed
+                    pct = (samples_written / max(1, total_estimated)) * 100
+                    
+                    mdb_path = os.path.join(output_dir, "data.mdb")
+                    db_mb = os.path.getsize(mdb_path) / (1024 * 1024) if os.path.exists(mdb_path) else 0
+                    
+                    sys.stdout.write(
+                        f"\r[*] [Synth Parallel] Đã ghi: {samples_written:,}/{total_estimated:,} ({pct:.1f}%) | "
+                        f"Tốc độ: {speed:,.0f} mẫu/s | LMDB: {db_mb:,.1f} MB  "
+                    )
+                    sys.stdout.flush()
+
+        # -------------------------------------------------------------
+        # PHASE 3: HuggingFace Parquet integration (if any specified)
+        # -------------------------------------------------------------
+        if hf_parquet_files:
+            print("\n[GIAI ĐOẠN 3] Nạp dữ liệu bổ sung từ Hugging Face Parquet...")
+            try:
+                import pandas as pd
+                for pq_file in hf_parquet_files:
+                    if os.path.exists(pq_file):
+                        df = pd.read_parquet(pq_file)
+                        for _, row in df.iterrows():
+                            # Find text and image bytes
+                            raw_bytes = None
+                            text = None
+                            for col in df.columns:
+                                if "text" in col.lower() or "label" in col.lower() or "transcription" in col.lower():
+                                    text = str(row[col]).strip()
+                                elif "image" in col.lower():
+                                    val = row[col]
+                                    if isinstance(val, dict) and "bytes" in val:
+                                        raw_bytes = val["bytes"]
+                                    elif isinstance(val, bytes):
+                                        raw_bytes = val
+                                        
+                            if raw_bytes and text:
+                                try:
+                                    img = Image.open(io.BytesIO(raw_bytes)).convert("RGB")
+                                    w, h = img.size
+                                    new_w = max(16, int((float(target_height) / max(1, h)) * w))
+                                    resized = img.resize((new_w, target_height), Image.LANCZOS)
+                                    buf = io.BytesIO()
+                                    resized.save(buf, format="JPEG", quality=90)
+                                    
+                                    samples_written += 1
+                                    img_key = f"image-{samples_written:09d}".encode("utf-8")
+                                    lbl_key = f"label-{samples_written:09d}".encode("utf-8")
+                                    src_key = f"source-{samples_written:09d}".encode("utf-8")
+                                    
+                                    txn.put(img_key, buf.getvalue())
+                                    txn.put(lbl_key, text.encode("utf-8"))
+                                    txn.put(src_key, b"hf_handwriting")
+                                    
+                                    stats["hf_handwriting"] = stats.get("hf_handwriting", 0) + 1
+                                    uncommitted += 1
+                                    if uncommitted >= commit_interval:
+                                        txn.commit()
+                                        txn = env.begin(write=True)
+                                        uncommitted = 0
+                                except Exception:
+                                    continue
+            except Exception as e:
+                print(f"[-] Lỗi nạp HuggingFace Parquet: {e}")
+
+        # Final commit and write metadata header
+        txn.put(b"num-samples", str(samples_written).encode("utf-8"))
+        txn.commit()
+        
+    finally:
+        env.sync()
+        env.close()
+
+    total_time = max(0.001, time.time() - start_time)
+    avg_speed = samples_written / total_time
+    mdb_size_mb = os.path.getsize(os.path.join(output_dir, "data.mdb")) / (1024 * 1024)
+    
+    # Save comprehensive summary JSON
+    summary = {
+        "dataset_name": "AsterOCR Unified Multi-Source Dataset",
+        "output_dir": output_dir,
+        "total_samples": samples_written,
+        "target_height": target_height,
+        "distribution": stats,
+        "size_mb": round(mdb_size_mb, 2),
+        "total_seconds": round(total_time, 2),
+        "avg_samples_per_sec": round(avg_speed, 1),
+        "created_at": time.strftime("%Y-%m-%d %H:%M:%S")
+    }
+    
+    with open(os.path.join(output_dir, "dataset_summary.json"), "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+
+    print("\n\n" + "=" * 72)
+    print("   [+] HOÀN THÀNH ĐÓNG GÓI BỘ DỮ LIỆU HỢP NHẤT DUY NHẤT!")
+    print("=" * 72)
+    print(f"  Tổng số mẫu trong kho : {samples_written:,} mẫu")
+    print(f"  Thời gian hoàn thành  : {total_time:.1f}s (~{total_time/60:.1f} phút)")
+    print(f"  Tốc độ ghi trung bình : {avg_speed:,.0f} mẫu/giây")
+    print(f"  Dung lượng LMDB (MDB) : {mdb_size_mb:,.2f} MB")
+    print(f"  Phân bố nguồn dữ liệu :")
+    for k, v in stats.items():
+        if v > 0:
+            pct = (v / max(1, samples_written)) * 100
+            print(f"    * {k:<20}: {v:>10,} mẫu ({pct:.1f}%)")
+    print("=" * 72)
+    return summary
+
+def main():
+    parser = argparse.ArgumentParser(description="AsterOCR Unified Multi-Source Dataset Packer")
+    parser.add_argument("--output_dir", type=str, default="dataset_unified.lmdb", help="Thư mục LMDB duy nhất")
+    parser.add_argument("--num_synth_hw", type=int, default=50000, help="Số mẫu chữ viết tay giả lập")
+    parser.add_argument("--num_synth_rc", type=int, default=25000, help="Số mẫu hóa đơn biên lai giả lập")
+    parser.add_argument("--max_real_hw", type=int, default=None, help="Giới hạn số mẫu viết tay thật (None = lấy hết 7.2k)")
+    parser.add_argument("--max_real_rc", type=int, default=50000, help="Giới hạn số mẫu hóa đơn thật")
+    parser.add_argument("--max_real_doc", type=int, default=50000, help="Giới hạn số mẫu văn bản tài liệu thật")
+    parser.add_argument("--no_real_zip", action="store_true", help="Không lấy dữ liệu từ zip")
+    parser.add_argument("--num_workers", type=int, default=10, help="Số luồng CPU tối đa (mặc định 10)")
+    parser.add_argument("--target_height", type=int, default=32, help="Chiều cao ảnh chuẩn hóa")
+    
+    args = parser.parse_args()
+    
+    pack_unified_dataset(
+        output_dir=args.output_dir,
+        num_synth_handwriting=args.num_synth_hw,
+        num_synth_receipt=args.num_synth_rc,
+        include_real_zip=not args.no_real_zip,
+        max_real_handwriting=args.max_real_hw,
+        max_real_receipts=args.max_real_rc,
+        max_real_documents=args.max_real_doc,
+        num_workers=args.num_workers,
+        target_height=args.target_height
+    )
+
+if __name__ == "__main__":
+    main()

asterocr/download_dataset.py

@@ -0,0 +1,93 @@
+import os
+import re
+import sys
+import requests
+
+if sys.platform == "win32":
+    try:
+        sys.stdout.reconfigure(encoding='utf-8')
+        sys.stderr.reconfigure(encoding='utf-8')
+    except Exception:
+        pass
+
+def get_filename_from_cd(cd):
+    if not cd:
+        return None
+    fname = re.findall(r'filename=(.+)', cd)
+    if len(fname) == 0:
+        fname = re.findall(r"filename\*=UTF-8''(.+)", cd)
+    if len(fname) > 0:
+        return fname[0].strip(' "\'')
+    return None
+
+def download_file_from_google_drive(file_id, output_dir="dataset_download"):
+    os.makedirs(output_dir, exist_ok=True)
+    session = requests.Session()
+    
+    # Initial request
+    url = f"https://drive.usercontent.google.com/download?id={file_id}&export=download&authuser=0"
+    headers = {
+        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
+    }
+    
+    print(f"[*] Đang kết nối tới Google Drive (ID: {file_id})...")
+    res = session.get(url, headers=headers, stream=True)
+    
+    content_type = res.headers.get('Content-Type', '')
+    if 'text/html' in content_type:
+        html = res.text
+        # Tìm token xác nhận hoặc link download trực tiếp
+        token_match = re.search(r'name="confirm"\s+value="([^"]+)"', html)
+        uuid_match = re.search(r'name="uuid"\s+value="([^"]+)"', html)
+        action_match = re.search(r'action="([^"]+)"', html)
+        
+        if action_match and token_match:
+            download_url = action_match.group(1)
+            params = {
+                'id': file_id,
+                'export': 'download',
+                'confirm': token_match.group(1),
+            }
+            if uuid_match:
+                params['uuid'] = uuid_match.group(1)
+            print(f"[*] Bỏ qua cảnh báo virus, gửi mã xác nhận: {token_match.group(1)}...")
+            res = session.get(download_url, params=params, headers=headers, stream=True)
+        else:
+            # Thử qua docs.google.com/uc
+            alt_url = "https://docs.google.com/uc?export=download"
+            res = session.get(alt_url, params={'id': file_id, 'confirm': 't'}, headers=headers, stream=True)
+
+    # Lấy tên file thực tế từ header hoặc đặt mặc định
+    filename = get_filename_from_cd(res.headers.get('Content-Disposition'))
+    if not filename:
+        filename = f"dataset_{file_id}.zip"
+        
+    out_path = os.path.join(output_dir, filename)
+    print(f"[*] Đang tải về: {filename}")
+    total_size = int(res.headers.get('content-length', 0))
+    if total_size:
+        print(f"[*] Dung lượng: {total_size / (1024*1024):.2f} MB")
+        
+    downloaded = 0
+    chunk_size = 1024 * 1024 # 1MB
+    with open(out_path, "wb") as f:
+        for chunk in res.iter_content(chunk_size):
+            if chunk:
+                f.write(chunk)
+                downloaded += len(chunk)
+                if total_size > 0:
+                    percent = (downloaded / total_size) * 100
+                    mb_down = downloaded / (1024 * 1024)
+                    mb_tot = total_size / (1024 * 1024)
+                    sys.stdout.write(f"\r--> Tiến độ: {percent:.1f}% ({mb_down:.1f}/{mb_tot:.1f} MB)")
+                else:
+                    mb_down = downloaded / (1024 * 1024)
+                    sys.stdout.write(f"\r--> Đã tải: {mb_down:.1f} MB")
+                sys.stdout.flush()
+                
+    print(f"\n[+] Tải xong thành công: {out_path}")
+    return out_path
+
+if __name__ == "__main__":
+    file_id = "1T0cmkhTgu3ahyMIwGZeby612RpVdDxOR"
+    download_file_from_google_drive(file_id, "dataset_download")