Feature: Integrated CNN model and optimized TrOCR with Beam Search

app.py

@@ -1,21 +1,24 @@
 import os
 import gradio as gr
 import torch
+import numpy as np
 from PIL import Image
 from peft import PeftModel
 from transformers import AutoTokenizer, TrOCRProcessor, ViTImageProcessor, VisionEncoderDecoderModel
+from cnn_model import CharacterClassifier  # Importing your CNN logic
 
 # --- CONFIGURATION ---
 BASE_MODEL_ID = "paudelanil/trocr-devanagari-2"
-ADAPTER_ID = "manishw10/devgen-trocr-devanagari-lora" 
+ADAPTER_ID = "manishw10/devgen-trocr-devanagari-lora"
+CNN_MODEL_PATH = "devanagari-cnn-classifier.pt"
 
 # Detect environment
 IS_SPACE = "SPACE_ID" in os.environ
-
-print(f"System: Loading model... (Env: {'Hugging Face Space' if IS_SPACE else 'Local'})")
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load Processor
+print(f"System: Initializing Models (Env: {'Hugging Face Space' if IS_SPACE else 'Local'})")
+
+# 1. Load TrOCR Model & Processor
 try:
     processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
 except Exception:
@@ -23,48 +26,82 @@ except Exception:
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
     processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer)
 
-# Load Model
 base_model = VisionEncoderDecoderModel.from_pretrained(BASE_MODEL_ID)
-# The PeftModel wrapper injects weights into base_model
 model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
 model.to(device)
 model.eval()
-print(f"System: Model loaded successfully on {device}")
 
-def predict(image):
+# 2. Load CNN Classifier
+cnn_engine = CharacterClassifier(model_path=CNN_MODEL_PATH, device=device)
+
+print(f"System: Models loaded successfully on {device}")
+
+def predict_trocr(image):
     if image is None:
         return "Error: No image uploaded"
     try:
         image = image.convert("RGB")
         pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
         
-        # --- THE ROBUST FIX ---
-        # We call .base_model.generate() directly.
-        # This bypasses the buggy PEFT wrapper while still using the LoRA weights.
-        # We also add max_new_tokens for a better result.
+        # --- HIGH-QUALITY GENERATION ---
+        # Added num_beams and length_penalty to fix the "rubbish" output.
+        # This makes TrOCR use Beam Search instead of Greedy Search.
         with torch.no_grad():
             generated_ids = model.base_model.generate(
                 pixel_values=pixel_values,
-                max_new_tokens=64
+                num_beams=4,
+                length_penalty=1.0,
+                max_new_tokens=64,
+                early_stopping=True,
+                decoder_start_token_id=model.config.decoder_start_token_id
             )
             generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         return generated_text
     except Exception as e:
         import traceback
-        print(traceback.format_exc()) # Log the full error to the Space logs
-        return f"Error during inference: {str(e)}"
+        print(traceback.format_exc())
+        return f"TrOCR Error: {str(e)}"
+
+def predict_cnn(image):
+    if image is None:
+        return "Error: No image uploaded"
+    try:
+        image = image.convert("RGB")
+        result = cnn_engine.predict(image)
+        if "error" in result:
+            return result["error"]
+        return f"Character: {result['text']} (Confidence: {result['confidence']:.2%})"
+    except Exception as e:
+        return f"CNN Error: {str(e)}"
+
+# --- CUSTOM GRADIO INTERFACE ---
+with gr.Blocks(title="DevGen OCR Suite") as demo:
+    gr.Markdown("# 🕉️ DevGen Devanagari OCR Suite")
+    gr.Markdown("Switch between TrOCR (for words/sentences) and CNN (for single characters).")
+    
+    with gr.Tabs():
+        with gr.TabItem("TrOCR (Word/Sentence Recognition)"):
+            with gr.Row():
+                with gr.Column():
+                    img_input = gr.Image(type="pil", label="Upload Handwritten Word")
+                    btn_trocr = gr.Button("Recognize Word", variant="primary")
+                with gr.Column():
+                    text_output = gr.Textbox(label="Recognized Text")
+            btn_trocr.click(fn=predict_trocr, inputs=img_input, outputs=text_output)
+            
+        with gr.TabItem("CNN (Single Character Recognition)"):
+            with gr.Row():
+                with gr.Column():
+                    char_input = gr.Image(type="pil", label="Upload Single Character")
+                    btn_cnn = gr.Button("Classify Character", variant="primary")
+                with gr.Column():
+                    char_output = gr.Textbox(label="Classification Result")
+            btn_cnn.click(fn=predict_cnn, inputs=char_input, outputs=char_output)
 
-# Interface setup
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload Handwritten Devanagari Word"),
-    outputs=gr.Textbox(label="Recognized Text"),
-    title="DevGen Devanagari OCR",
-    description="Recognize handwritten Devanagari words using TrOCR and LoRA adaptation.",
-    allow_flagging="never"
-)
+    gr.Markdown("---")
+    gr.Markdown("Built with ❤️ by DevGen Team. Using TrOCR + LoRA and custom 3-layer CNN.")
 
 if __name__ == "__main__":
-    # If running on HF Spaces, use 0.0.0.0. If local, use default localhost.
     server_name = "0.0.0.0" if IS_SPACE else "127.0.0.1"
+    # Note: We don't use monkey-patching here, the base_model.generate handles it.
     demo.launch(server_name=server_name)

cnn_model.py

@@ -0,0 +1,226 @@
+"""
+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

devanagari-cnn-classifier.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b78f582dabc34056303acf582b5b97a47ac32fa847bfeb33227c1a41681f760
+size 2719832

requirements.txt

@@ -7,3 +7,6 @@ safetensors
 sentencepiece
 fastapi<0.113.0
 uvicorn
+opencv-python-headless
+torchvision
+numpy