asterocr 1.0.0__py3-none-any.whl

asterocr/dataset/loader.py

@@ -0,0 +1,371 @@
+"""
+AsterOCR Multi-Source Blended Dataset Engine
+Features:
+1. AsterZipDataset: Reads directly from 7.3GB ocr_dataset.zip without extracting to disk.
+   Caches sample index to zip_index.json for instant 0.05s loading.
+2. AsterSyntheticDataset: Infinite realistic Vietnamese administrative text generator.
+3. AsterMultiSourceDataset: Blends real-world scanned images + synthetic Vietnamese data
+   according to configurable mixing ratios (e.g. 50% Real / 50% Synthetic).
+4. ctc_collate_fn: Batch padder and CTC length tensor builder.
+"""
+
+import os
+import io
+import json
+import random
+import zipfile
+import torch
+from torch.utils.data import Dataset, DataLoader
+from PIL import Image
+import numpy as np
+
+try:
+    import lmdb
+except ImportError:
+    lmdb = None
+
+from asterocr.vocab import default_vocab
+from asterocr.synthesizer.corpus import generate_random_corpus_text
+from asterocr.synthesizer.generator import default_generator
+
+class AsterLMDBDataset(Dataset):
+    """
+    Ultra-high performance zero-copy LMDB Dataset loader for AsterOCR.
+    Reads directly from LMDB database (e.g. dataset_mega.lmdb).
+    Compatible with standard CRNN / ASTER / Clova AI format:
+      - 'image-%09d' -> raw JPEG/PNG bytes
+      - 'label-%09d' -> UTF-8 encoded text string
+      - 'num-samples' -> total number of samples
+    Supports multi-processing PyTorch DataLoader via lazy environment initialization.
+    """
+    def __init__(self, lmdb_dir, target_height=32, max_samples=None):
+        if lmdb is None:
+            raise ImportError("Gói 'lmdb' chưa được cài đặt! Hãy chạy: pip install lmdb")
+            
+        self.lmdb_dir = lmdb_dir
+        self.target_height = target_height
+        self.env = None
+        
+        if not os.path.exists(lmdb_dir):
+            raise FileNotFoundError(f"Thư mục LMDB không tồn tại: {lmdb_dir}")
+            
+        from asterocr.kaggle_bridge.chunker import get_or_open_lmdb_env
+        self.env = get_or_open_lmdb_env(self.lmdb_dir, max_readers=126)
+        if self.env is None:
+            raise FileNotFoundError(f"Không thể mở LMDB tại: {lmdb_dir}")
+            
+        with self.env.begin(write=False) as txn:
+            num_samples_val = txn.get(b'num-samples')
+            if num_samples_val is None:
+                raise ValueError(f"Không tìm thấy khóa 'num-samples' trong LMDB: {lmdb_dir}")
+            total = int(num_samples_val.decode('utf-8'))
+            self.num_samples = min(total, max_samples) if (max_samples and max_samples > 0) else total
+
+    def _init_env(self):
+        if self.env is None:
+            from asterocr.kaggle_bridge.chunker import get_or_open_lmdb_env
+            self.env = get_or_open_lmdb_env(self.lmdb_dir, max_readers=126)
+
+    def __len__(self):
+        return self.num_samples
+
+    def __getitem__(self, idx):
+        self._init_env()
+        sample_idx = idx + 1
+        try:
+            with self.env.begin(write=False) as txn:
+                img_key = f"image-{sample_idx:09d}".encode("utf-8")
+                lbl_key = f"label-{sample_idx:09d}".encode("utf-8")
+                
+                img_bytes = txn.get(img_key)
+                label_bytes = txn.get(lbl_key)
+                
+                if img_bytes is None or label_bytes is None:
+                    img = Image.new("RGB", (64, self.target_height), (255, 255, 255))
+                    text = " "
+                else:
+                    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+                    raw_str = label_bytes.decode("utf-8", errors="ignore").strip()
+                    text = raw_str if raw_str else " "
+        except Exception:
+            img = Image.new("RGB", (64, self.target_height), (255, 255, 255))
+            text = " "
+
+        w, h = img.size
+        new_w = max(16, int((float(self.target_height) / max(1, h)) * w))
+        # Guarantee enough width for characters so CTC feature length >= target length
+        min_w = len(text) * 8 + 16
+        new_w = min(480, max(new_w, min_w))
+        img = img.resize((new_w, self.target_height), Image.BILINEAR)
+
+        arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
+        return torch.from_numpy(arr), text
+
+class AsterSyntheticDataset(Dataset):
+    """Infinite synthetic training dataset generated in real time."""
+    def __init__(self, size=10000, target_height=32, generator=None):
+        self.size = size
+        self.target_height = target_height
+        self.generator = generator or default_generator
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        text = generate_random_corpus_text()
+        img = self.generator.render_text(text, target_height=self.target_height)
+        
+        arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
+        tensor = torch.from_numpy(arr)
+        return tensor, text
+
+class AsterZipDataset(Dataset):
+    """
+    Directly streams image and text label pairs from ocr_dataset.zip.
+    Zero disk extraction required.
+    """
+    def __init__(self, zip_path="dataset_download/ocr_dataset.zip", target_height=32, max_samples=None):
+        self.zip_path = zip_path
+        self.target_height = target_height
+        self.pairs = [] # list of (img_name, txt_name)
+        
+        index_cache = os.path.splitext(zip_path)[0] + "_index.json"
+        
+        if os.path.exists(index_cache):
+            try:
+                with open(index_cache, "r", encoding="utf-8") as f:
+                    self.pairs = json.load(f)
+            except Exception:
+                self.pairs = []
+
+        if not self.pairs and os.path.exists(zip_path):
+            print(f"[*] Đang lập chỉ mục {zip_path} (lần đầu tiên)...")
+            with zipfile.ZipFile(zip_path, "r") as z:
+                all_names = set(z.namelist())
+                # Filter for images that have matching txt files
+                for n in all_names:
+                    if n.endswith(('.jpg', '.png')):
+                        base, _ = os.path.splitext(n)
+                        txt_n = base + ".txt"
+                        if txt_n in all_names:
+                            self.pairs.append((n, txt_n))
+            
+            # Save index cache
+            try:
+                with open(index_cache, "w", encoding="utf-8") as f:
+                    json.dump(self.pairs, f)
+                print(f"[+] Đã lưu chỉ mục: {len(self.pairs)} mẫu vào {index_cache}")
+            except Exception:
+                pass
+                
+        if max_samples and len(self.pairs) > max_samples:
+            self.pairs = self.pairs[:max_samples]
+            
+        self.zip_ref = None
+
+    def _get_zip(self):
+        if self.zip_ref is None:
+            self.zip_ref = zipfile.ZipFile(self.zip_path, "r")
+        return self.zip_ref
+
+    def __len__(self):
+        return len(self.pairs)
+
+    def __getitem__(self, idx):
+        if not self.pairs:
+            # Fallback dummy
+            img = Image.new("RGB", (64, self.target_height), (255, 255, 255))
+            arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
+            return torch.from_numpy(arr), " "
+
+        img_n, txt_n = self.pairs[idx]
+        z = self._get_zip()
+        
+        try:
+            raw_img = z.read(img_n)
+            img = Image.open(io.BytesIO(raw_img)).convert("RGB")
+            w, h = img.size
+            new_w = max(16, int((self.target_height / max(1, h)) * w))
+            img = img.resize((new_w, self.target_height), Image.BILINEAR)
+            
+            raw_txt = z.read(txt_n).decode("utf-8", errors="ignore").strip().strip('"\'')
+            text = raw_txt if raw_txt else " "
+        except Exception:
+            img = Image.new("RGB", (64, self.target_height), (255, 255, 255))
+            text = " "
+
+        arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
+        return torch.from_numpy(arr), text
+
+class AsterMultiSourceDataset(Dataset):
+    """
+    Multi-Source Blended Dataset:
+    Dynamically blends Real Scans/Photos (from zip) + Synthetic Vietnamese text
+    according to a configurable ratio (e.g. 50% Real / 50% Synthetic).
+    """
+    def __init__(self, synthetic_ds, real_ds=None, real_ratio=0.5, total_size=10000):
+        self.synthetic_ds = synthetic_ds
+        self.real_ds = real_ds
+        self.real_ratio = real_ratio if (real_ds and len(real_ds) > 0) else 0.0
+        self.total_size = total_size
+
+    def __len__(self):
+        return self.total_size
+
+    def __getitem__(self, idx):
+        # Decide whether to sample from real dataset or synthetic generator
+        if random.random() < self.real_ratio and self.real_ds:
+            real_idx = random.randint(0, len(self.real_ds) - 1)
+            return self.real_ds[real_idx]
+        else:
+            synth_idx = random.randint(0, len(self.synthetic_ds) - 1)
+            return self.synthetic_ds[synth_idx]
+
+class AsterOmniDataset(Dataset):
+    """
+    Omni-Source Blended Dataset:
+    Hop nhat toan bo moi nguon du lieu (All-In-One):
+    - Kho Hop Nhat LMDB (Viet tay VNOnDB, Vo o ly, Hoa don in kim)
+    - Kho Anh Quet/Chup Thuc Te (601,282 mau trong zip)
+    - Kho Mega Trieu Mau LMDB (1,000,000 mau sieu toc)
+    - Trinh Tu Sinh Dong Python On-The-Fly (322 Font Windows, Giay cu, Muc nhoe)
+    Can bang ty trong mau va boc ngau nhien truc tiep cho tung batch.
+    """
+    def __init__(self, sources_dict, weights=None, total_size=50000):
+        self.sources = []
+        self.names = []
+        for name, ds in sources_dict.items():
+            if ds is not None and len(ds) > 0:
+                self.sources.append(ds)
+                self.names.append(name)
+
+        if not self.sources:
+            raise ValueError("Khong tim thay nguon du lieu nao kha dung cho Omni Dataset!")
+
+        self.total_size = total_size
+        n = len(self.sources)
+        if weights is None or len(weights) != n:
+            self.weights = [1.0 / n] * n
+        else:
+            s = sum(weights)
+            self.weights = [w / s for w in weights]
+
+    def __len__(self):
+        return self.total_size
+
+    def __getitem__(self, idx):
+        chosen_ds = random.choices(self.sources, weights=self.weights, k=1)[0]
+        rand_idx = random.randint(0, len(chosen_ds) - 1)
+        return chosen_ds[rand_idx]
+
+class AsterFileDataset(Dataset):
+    """Loads dataset from image directory and annotation file (path\\tlabel)."""
+    def __init__(self, annotation_file, data_root=None, target_height=32):
+        self.samples = []
+        self.target_height = target_height
+        data_root = data_root or os.path.dirname(annotation_file)
+        
+        with open(annotation_file, "r", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                parts = line.split("\t")
+                if len(parts) >= 2:
+                    rel_p, text = parts[0], parts[1]
+                else:
+                    parts = line.split(" ", 1)
+                    rel_p = parts[0]
+                    text = parts[1] if len(parts) > 1 else ""
+                    
+                full_p = os.path.join(data_root, rel_p) if not os.path.isabs(rel_p) else rel_p
+                if os.path.exists(full_p):
+                    self.samples.append((full_p, text))
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        img_path, text = self.samples[idx]
+        try:
+            img = Image.open(img_path).convert("RGB")
+            w, h = img.size
+            new_w = max(16, int((self.target_height / max(1, h)) * w))
+            img = img.resize((new_w, self.target_height), Image.BILINEAR)
+        except Exception:
+            img = Image.new("RGB", (32, self.target_height), (255, 255, 255))
+            text = " "
+            
+        arr = np.array(img, dtype=np.float32).transpose(2, 0, 1) / 255.0
+        return torch.from_numpy(arr), text
+
+def ctc_collate_fn(batch, vocab=None):
+    if vocab is None:
+        vocab = default_vocab
+        
+    tensors, texts = zip(*batch)
+    batch_size = len(tensors)
+    
+    encoded_targets = []
+    target_lengths = []
+    
+    for text in texts:
+        tokens = vocab.encode(text)
+        if not tokens:
+            tokens = [vocab.encode(" ")[0]]
+        encoded_targets.extend(tokens)
+        target_lengths.append(len(tokens))
+        
+    targets = torch.tensor(encoded_targets, dtype=torch.long)
+    target_lengths = torch.tensor(target_lengths, dtype=torch.long)
+    max_target_len = int(target_lengths.max().item()) if len(target_lengths) > 0 else 1
+
+    max_w = min(640, max([t.shape[2] for t in tensors]))
+    max_w = max(max_w, (max_target_len * 4) + 16)
+    max_w = max(32, ((max_w + 3) // 4) * 4)
+    
+    padded_imgs = torch.ones((batch_size, 3, 32, max_w), dtype=torch.float32)
+    input_lengths = []
+    
+    for i, t in enumerate(tensors):
+        w = min(t.shape[2], max_w)
+        padded_imgs[i, :, :, :w] = t[:, :, :w]
+        feat_len = max(1, w // 4)
+        t_len = int(target_lengths[i].item())
+        feat_len = max(feat_len, t_len)
+        feat_len = min(feat_len, max_w // 4)
+        input_lengths.append(feat_len)
+        
+    input_lengths = torch.tensor(input_lengths, dtype=torch.long)
+    
+    return {
+        "images": padded_imgs,
+        "targets": targets,
+        "input_lengths": input_lengths,
+        "target_lengths": target_lengths,
+        "texts": list(texts)
+    }
+
+def create_dataloader(dataset, batch_size=32, shuffle=True, num_workers=0, pin_memory=False):
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        collate_fn=ctc_collate_fn,
+        pin_memory=pin_memory
+    )
+
+if __name__ == "__main__":
+    # Test Multi-Source Mixed Dataset
+    synth_ds = AsterSyntheticDataset(size=50)
+    zip_p = "dataset_download/ocr_dataset.zip"
+    real_ds = AsterZipDataset(zip_p) if os.path.exists(zip_p) else None
+    mixed_ds = AsterMultiSourceDataset(synth_ds, real_ds, real_ratio=0.5, total_size=20)
+    loader = create_dataloader(mixed_ds, batch_size=4)
+    
+    for batch in loader:
+        print("Mixed batch images:", batch["images"].shape)
+        print("Mixed batch texts:", batch["texts"])
+        break
+    print("[+] Multi-source dataset test passed!")