--- language: - km - en base_model: tharas/resnet_bilstm_ctc_100k pipeline_tag: image-to-text tags: - ocr - khmer - crnn - resnet - ctc --- # Khmer OCR — ResNet + BiLSTM + CTC This repository contains a deep learning model designed for Khmer and English Optical Character Recognition (OCR). It utilizes a ResNet backbone for spatial feature extraction, a bidirectional LSTM for sequence modeling, and Connectionist Temporal Classification (CTC) loss for alignment-free text recognition. ## Model Details ### Model Description - **Model type:** `khm_ocr_general_document` - **Language(s):** Khmer (`km`), English (`en`). *Note: The model exhibits higher accuracy and optimization patterns for Khmer character compositions.* - **Training State:** Trained from scratch on a comprehensive dataset of printed text lines. - **Architecture:** - **CNN:** ResNet blocks processing grayscale document line images downscaled to a fixed height. - **RNN:** 2-layer Bidirectional LSTM tracking character transitions. - **Classifier:** Linear layer mapping features to vocabulary indices decoded via greedy CTC. - **Charactor Error Rate:** 0.005589 - **Word Error Rate:** 0.045868 - **Training**: - **batch_size:** 16 - **device:** "cuda" - **epochs:** 30 - **learning_rate:** 0.001 - **num_workers:** 4 - **optimizer:** "adam" - **scheduler_factor:** 0.5 - **scheduler_patience:** 5 ### Out-of-Scope Use While this model handles varied font distributions well, it is strictly an line-level text recognizer. * **Optimal on:** Printed text documents featuring standard degradation, light blur, or low levels of scanner noise. * **Fails on:** Heavily underlined text paths, highly blurred captures, handwritten manuscripts, or unwarped, rotated document images. Ensure input crops are straight and baseline-aligned before inference. ### Recommendations This architecture is under active development. While it successfully handles clean, isolated line crops, performance drops on extreme layouts. Preprocessing elements (like layout analysis and text-line deskewing) should be executed prior to feeding imagery to this network. --- ## How to Get Started with the Model ### File Layout Requirements Before running predictions, ensure your repository or local folder contains your model files and characters map structured as follows: ```text . ├── predict.py └── khmer_ocr_model_CRNN/ ├── best_model.pth └── char.json ``` ```predict.py``` ```bash import argparse import json from pathlib import Path import torch import torch.nn as nn from PIL import Image import torchvision.transforms as T HERE = Path(__file__).parent IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp"} # ── Model Architecture ──────────────────────────────────────────────────────── class ResBlock(nn.Module): def __init__(self, in_ch, out_ch, stride=1): super().__init__() self.conv = nn.Sequential( nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True), nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), ) self.downsample = None if stride != 1 or in_ch != out_ch: self.downsample = nn.Sequential( nn.Conv2d(in_ch, out_ch, 1, stride=stride, bias=False), nn.BatchNorm2d(out_ch), ) self.relu = nn.ReLU(inplace=True) def forward(self, x): identity = x out = self.conv(x) if self.downsample is not None: identity = self.downsample(x) return self.relu(out + identity) class ResNetBiLSTMCTC(nn.Module): def __init__(self, num_classes, hidden_size=256, num_layers=2, dropout=0.3): super().__init__() self.cnn = nn.Sequential( nn.Conv2d(1, 64, 3, padding=1, bias=False), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), nn.Sequential(ResBlock(64, 64, stride=1)), nn.Sequential(ResBlock(64, 128, stride=2)), nn.Sequential(ResBlock(128, 256, stride=1)), ) self.rnn = nn.LSTM( input_size=256 * 16, hidden_size=hidden_size, num_layers=num_layers, bidirectional=True, dropout=dropout if num_layers > 1 else 0.0, batch_first=False, ) self.classifier = nn.Linear(hidden_size * 2, num_classes) def forward(self, x): feat = self.cnn(x) # [B, 256, H', W'] b, c, h, w = feat.shape seq = feat.permute(3, 0, 1, 2) # [W', B, C, H'] seq = seq.reshape(w, b, c * h) # [W', B, C*H'] out, _ = self.rnn(seq) # [W', B, 2*hidden] return self.classifier(out).log_softmax(2) # [W', B, num_classes] # ── Charset Parsing ─────────────────────────────────────────────────────────── def load_charset(path): """ Builds an index-to-character mapping lookup from char.json files. Accommodates categories grouped by blocks, lists, or custom string token pairs. """ with open(path, encoding="utf-8") as f: data = json.load(f) if isinstance(data, list): return {i: ch for i, ch in enumerate(data)} if isinstance(data, dict): known_categories = {"khmer", "latin", "digits", "special"} if known_categories & set(data.keys()): order = ["khmer", "latin", "digits", "special"] all_chars = "".join(data.get(cat, "") for cat in order) return {0: " ", **{i + 1: ch for i, ch in enumerate(all_chars)}} try: return {int(k): v for k, v in data.items()} except ValueError: pass return {v: k for k, v in data.items()} raise ValueError(f"Unrecognized token map schema inside character file: {path}") # ── Processing & Greedy Decoding ───────────────────────────────────────────── def ctc_decode(log_probs, idx2char, blank=0): """Collapses consecutive duplicate indexes and strips blank tokens.""" indices = log_probs.argmax(dim=1).tolist() chars, prev = [], None for idx in indices: if idx != prev and idx != blank: chars.append(idx2char.get(idx, "?")) prev = idx return "".join(chars) _transform = T.Compose([ T.ToTensor(), T.Normalize(mean=[0.5], std=[0.5]), ]) def preprocess(image_path, img_h=64, img_w=512): """Converts target image asset to grayscale, scales to expected shape, and adds batch dims.""" img = Image.open(image_path).convert("L") img = img.resize((img_w, img_h), Image.BICUBIC) return _transform(img).unsqueeze(0) def predict_one(model, image_path, idx2char, device, blank=0): tensor = preprocess(image_path).to(device) with torch.no_grad(): log_probs = model(tensor)[:, 0, :] # Extracted time steps: [Time, Classes] return ctc_decode(log_probs.cpu(), idx2char, blank=blank) # ── Core Runtime Entrypoint ─────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser(description="Khmer OCR Inference System Stack") parser.add_argument( "input", help="Target filepath to standalone line crop OR parent path containing image lists.", ) parser.add_argument( "--model", default=str(HERE / "khmer_ocr_model_CRNN" / "best_model.pth"), help="Checkpoint parameter file destination location path.", ) parser.add_argument( "--charset", default=str(HERE / "khmer_ocr_model_CRNN" / "char.json"), help="JSON configuration text format vocabulary parsing schema.", ) parser.add_argument( "--output", default=str(HERE / "predictions.txt"), help="Target text document destination to record string output arrays.", ) parser.add_argument( "--device", default="cuda" if torch.cuda.is_available() else "cpu", help="Hardware execution pipeline override runtime flag.", ) args = parser.parse_args() device = torch.device(args.device) # Initialize vocabulary bounds idx2char = load_charset(args.charset) num_classes = max(idx2char.keys()) + 1 # Instantiate weights mapping sequence layout model = ResNetBiLSTMCTC(num_classes=num_classes) state_dict = torch.load(args.model, map_location="cpu") model.load_state_dict(state_dict) model.to(device).eval() print(f"Model File : {args.model}") print(f"Charset File : {args.charset} ({num_classes} distribution classes)") print(f"Target Device: {device}\n") input_path = Path(args.input) out_path = Path(args.output) # ── Path Evaluator Logic: Independent Image Evaluation ────────────────── if input_path.is_file(): text = predict_one(model, input_path, idx2char, device) print(f"Image File : {input_path.name}") print(f"Prediction : {text}") with open(out_path, "w", encoding="utf-8") as f: f.write(f"{input_path.name}\t{text}\n") print(f"Saved logs : {out_path}") return # ── Path Evaluator Logic: Directory Iteration Loop ─────────────────────── if input_path.is_dir(): images = sorted( p for p in input_path.iterdir() if p.suffix.lower() in IMAGE_EXTS ) if not images: print(f"Termination: No valid image file variations found inside directory context '{input_path}'") return print(f"Queued Processing Run: Found {len(images)} sequence elements inside directory structure.\n") with open(out_path, "w", encoding="utf-8") as f: for i, img_path in enumerate(images, 1): try: text = predict_one(model, img_path, idx2char, device) except Exception as error_exception: text = f"RUNTIME ERROR METRIC EXCEPTION: {error_exception}" print(f"[{i}/{len(images)}] {img_path.name} -> {text}") f.write(f"{img_path.name}\t{text}\n") print(f"\nExecution Pipeline Complete: Stream saved down cleanly into '{out_path}'") return print(f"Invalid Operation Error: Resource tracking index target reference '{input_path}' does not point to structural nodes.") if __name__ == "__main__": main() ```
