""" DevGen Framework — CNN Devanagari Character Classifier A lightweight CNN for classifying individual handwritten Devanagari characters (vowels, consonants, digits) — 46 classes total. This model complements TrOCR (which handles words) by handling single characters that TrOCR hallucinates on. """ from __future__ import annotations import os import time from pathlib import Path from typing import Optional import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from torchvision import transforms # ── 46-class label map ────────────────────────────────────────────────────── # Standard DHCD ordering: 36 consonants/vowels + 10 digits DEVANAGARI_CLASSES = [ # Consonants (ka to gya) "क", "ख", "ग", "घ", "ङ", "च", "छ", "ज", "झ", "ञ", "ट", "ठ", "ड", "ढ", "ण", "त", "थ", "द", "ध", "न", "प", "फ", "ब", "भ", "म", "य", "र", "ल", "व", "श", "ष", "स", "ह", "क्ष", "त्र", "ज्ञ", # Digits (0-9) "०", "१", "२", "३", "४", "५", "६", "७", "८", "९", ] # Reverse map: character → index CHAR_TO_INDEX = {ch: i for i, ch in enumerate(DEVANAGARI_CLASSES)} NUM_CLASSES = len(DEVANAGARI_CLASSES) # Default model path DEFAULT_CNN_MODEL_PATH = "devanagari-cnn-classifier.pt" class DevanagariCNN(nn.Module): """ 3-layer CNN for 32×32 grayscale character images. ~500K parameters — fast inference even on CPU. """ def __init__(self, num_classes: int = NUM_CLASSES): super().__init__() self.features = nn.Sequential( # Block 1: 32×32 → 16×16 nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.Conv2d(32, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Dropout2d(0.25), # Block 2: 16×16 → 8×8 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.Conv2d(64, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Dropout2d(0.25), # Block 3: 8×8 → 4×4 nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.AdaptiveAvgPool2d(4), nn.Dropout2d(0.25), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(128 * 4 * 4, 256), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(256, num_classes), ) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.features(x) x = self.classifier(x) return x # ── Inference transform ───────────────────────────────────────────────────── # Matches training: resize to 32×32, grayscale, normalize INFERENCE_TRANSFORM = transforms.Compose([ transforms.Grayscale(num_output_channels=1), transforms.Resize((32, 32)), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]), ]) class CharacterClassifier: """ Wrapper for loading and running the trained CNN model. Used by the smart router in TrOCREngine. """ def __init__( self, model_path: Optional[str] = None, device: Optional[str] = None, ): self.device = device or ("mps" if torch.backends.mps.is_available() else "cpu") # Find model file if model_path is None: project_root = Path(__file__).resolve().parent.parent model_path = str(project_root / DEFAULT_CNN_MODEL_PATH) self.model_path = model_path self.model: Optional[DevanagariCNN] = None self.available = False if os.path.exists(model_path): self._load_model() else: print(f"[CNN Classifier] Model not found at {model_path} — single character recognition disabled") def _load_model(self): """Load the trained CNN weights.""" try: self.model = DevanagariCNN(NUM_CLASSES) state_dict = torch.load(self.model_path, map_location=self.device, weights_only=True) self.model.load_state_dict(state_dict) self.model.to(self.device) self.model.eval() self.available = True size_mb = os.path.getsize(self.model_path) / 1e6 print(f"[CNN Classifier] Loaded ({size_mb:.1f} MB) on {self.device} — {NUM_CLASSES} classes") except Exception as exc: print(f"[CNN Classifier] Failed to load model: {exc}") self.model = None self.available = False def predict(self, image: Image.Image) -> dict: """ Classify a single character image. Returns: dict with text, confidence, class_index, model_used """ if not self.available or self.model is None: return {"text": "", "confidence": 0.0, "error": "CNN model not loaded"} started_at = time.perf_counter() # Preprocess using DHCD style tensor = self._preprocess_dhcd_style(image).unsqueeze(0).to(self.device) with torch.inference_mode(): logits = self.model(tensor) probs = F.softmax(logits, dim=1) confidence, pred_idx = probs.max(dim=1) predicted_char = DEVANAGARI_CLASSES[pred_idx.item()] conf_value = round(confidence.item(), 4) inference_ms = round((time.perf_counter() - started_at) * 1000, 2) # Top-3 predictions for debugging top3_probs, top3_indices = probs.topk(3, dim=1) top3 = [ {"char": DEVANAGARI_CLASSES[idx.item()], "confidence": round(prob.item(), 4)} for idx, prob in zip(top3_indices[0], top3_probs[0]) ] return { "text": predicted_char, "confidence": conf_value, "class_index": pred_idx.item(), "top3": top3, "inference_ms": inference_ms, "model_used": "cnn_classifier", } def _preprocess_dhcd_style(self, image: Image.Image) -> torch.Tensor: """Preprocesses a character image to match DHCD dataset (inverted, tightly cropped, padded).""" import cv2 import numpy as np # Convert PIL to CV2 grayscale img = np.array(image.convert("L")) # Binarize and invert (DHCD is white ink on black background) _, binary = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU) # Crop to bounding box contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) if contours: c = max(contours, key=cv2.contourArea) x, y, w, h = cv2.boundingRect(c) cropped = binary[y:y+h, x:x+w] else: cropped = binary h, w = cropped.shape # Pad to square and add 16px border (helps CNN focus on center) side = max(w, h) padded = np.zeros((side + 16, side + 16), dtype=np.uint8) y_off = (side + 16 - h) // 2 x_off = (side + 16 - w) // 2 padded[y_off:y_off+h, x_off:x_off+w] = cropped # Resize to 32x32 resized = cv2.resize(padded, (32, 32), interpolation=cv2.INTER_AREA) # Convert to tensor and normalize to [-1, 1] tensor = torch.tensor(resized, dtype=torch.float32).unsqueeze(0) tensor = (tensor / 255.0 - 0.5) / 0.5 return tensor