You can deploy this model and start earning money today!

Thareah/
khmer_ocr_model_CRNN

Name: khmer_ocr_model_CRNN
Availability: InStock

Views

No views yet

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: "<blank>", **{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()

Thareah/khmer_ocr_model_CRNN