Final Fidelity: Mirror local Preprocess-Route-Recognize pipeline

app.py

@@ -1,67 +1,55 @@
 import os
-import time
+import io
 import gradio as gr
 import torch
 import numpy as np
-import cv2
 from PIL import Image
 from peft import PeftModel
-from transformers import AutoTokenizer, TrOCRProcessor, ViTImageProcessor, VisionEncoderDecoderModel
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from cnn_model import CharacterClassifier
+from preprocessing import preprocess_for_ocr
 
-# --- CONSTANTS (From local trocr_engine.py) ---
+# --- CONFIGURATION ---
 BASE_MODEL_ID = "paudelanil/trocr-devanagari-2"
 ADAPTER_ID = "manishw10/devgen-trocr-devanagari-lora"
 CNN_MODEL_PATH = "devanagari-cnn-classifier.pt"
-SPECIAL_TOKEN_NAMES = ("bos_token_id", "cls_token_id", "eos_token_id", "pad_token_id", "sep_token_id")
 
-IS_SPACE = "SPACE_ID" in os.environ
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# --- MODEL LOADING (Mirrored from TrOCREngine.__init__) ---
-print(f"System: Aligning with local TrOCREngine...")
-
-# 1. Load Processor
-try:
-    processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
-except Exception:
-    image_processor = ViTImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
-    processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer)
-
-# 2. Load and Config Model
+# --- MODEL INITIALIZATION (Mirror of TrOCREngine) ---
+print("System: Initializing Full-Fidelity Engine...")
+processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
 base_model = VisionEncoderDecoderModel.from_pretrained(BASE_MODEL_ID)
+
+# Sync Token Configs
 base_model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
 base_model.config.pad_token_id = processor.tokenizer.pad_token_id
 base_model.config.eos_token_id = processor.tokenizer.sep_token_id
 base_model.config.vocab_size = base_model.config.decoder.vocab_size
 
-# 3. Apply and MERGE LoRA (Critical for consistency)
+# Apply and Merge PEFT (Identical to local)
 peft_model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
 try:
     model = peft_model.merge_and_unload()
-    print("System: LoRA weights merged successfully (1:1 with local).")
-except Exception as e:
-    print(f"System: Merge failed, using wrapper: {e}")
+    print("System: LoRA weights merged.")
+except Exception:
     model = peft_model
-
 model.to(device)
 model.eval()
 
-# 4. Load CNN
+# Load CNN
 cnn_engine = CharacterClassifier(model_path=CNN_MODEL_PATH, device=device)
 
-# --- ORIGINAL ROUTING (Mirrored from image_router.py) ---
+# --- ORIGINAL ROUTING LOGIC (1:1 Copy) ---
 def _flood_fill(binary, visited, start_y, start_x, h, w):
     stack = [(start_y, start_x)]
     size = 0
     while stack:
         y, x = stack.pop()
-        if y < 0 or y >= h or x < 0 or x >= w or visited[y, x] or not binary[y, x]:
-            continue
-        visited[y, x] = True
+        if y<0 or y>=h or x<0 or x>=w or visited[y,x] or not binary[y,x]: continue
+        visited[y,x] = True
         size += 1
-        stack.extend([(y+1, x), (y-1, x), (y, x+1), (y, x-1)])
+        stack.extend([(y+1,x),(y-1,x),(y,x+1),(y,x-1)])
     return size
 
 def count_blobs(binary, min_size=10):
@@ -70,91 +58,75 @@ def count_blobs(binary, min_size=10):
     count = 0
     for y in range(h):
         for x in range(w):
-            if binary[y, x] and not visited[y, x]:
+            if binary[y,x] and not visited[y,x]:
                 size = _flood_fill(binary, visited, y, x, h, w)
-                if size >= min_size:
-                    count += 1
+                if size >= min_size: count += 1
     return count
 
-def classify_input_type(image):
+def original_classify_input(image):
     gray = image.convert("L")
     arr = np.array(gray)
     threshold = min(arr.mean() * 0.75, 200)
     binary = (arr < threshold).astype(np.uint8)
-    
-    rows = np.any(binary, axis=1)
-    cols = np.any(binary, axis=0)
-    if not rows.any() or not cols.any():
-        return "character", 0.5, 0.0, 0
-    
+    rows, cols = np.any(binary, axis=1), np.any(binary, axis=0)
+    if not rows.any() or not cols.any(): return "character", 0, 0
     rmin, rmax = np.where(rows)[0][[0, -1]]
     cmin, cmax = np.where(cols)[0][[0, -1]]
     w, h = cmax - cmin + 1, rmax - rmin + 1
-    aspect_ratio = w / max(h, 1)
-    blob_count = count_blobs(binary, min_size=max(binary.size * 0.001, 10))
+    ar, bc = w/h, count_blobs(binary, min_size=max(binary.size * 0.001, 10))
+    is_char = True
+    if ar > 2.5: is_char = False
+    elif ar > 1.8 and bc >= 3: is_char = False
+    elif bc >= 4: is_char = False
+    elif ar < 1.3 and bc <= 2: is_char = True
+    elif bc == 1 and ar < 1.5: is_char = True
+    elif ar > 1.6: is_char = False
+    return ("character" if is_char else "word"), ar, bc
+
+# --- MAIN INFERENCE PIPELINE ---
+def predict(image):
+    if image is None: return None, "Upload image.", ""
     
-    # DECISION LOGIC (1:1 with local image_router.py)
-    is_character = True
-    if aspect_ratio > 2.5: is_character = False
-    elif aspect_ratio > 1.8 and blob_count >= 3: is_character = False
-    elif blob_count >= 4: is_character = False
-    elif aspect_ratio < 1.3 and blob_count <= 2: is_character = True
-    elif blob_count == 1 and aspect_ratio < 1.5: is_character = True
-    elif aspect_ratio > 1.6: is_character = False
+    # 1. PREPROCESS (Critical! 1:1 with local recognize endpoint)
+    # Convert PIL to bytes for the preprocessor
+    buf = io.BytesIO()
+    image.save(buf, format="PNG")
+    image_bytes = buf.getvalue()
     
-    return ("character" if is_character else "word"), aspect_ratio, blob_count
-
-# --- PREDICT ---
-def predict(image, manual_mode):
-    if image is None: return None, "Upload an image.", "", ""
+    # The original app preprocesses BEFORE routing
+    preprocessed_pil = preprocess_for_ocr(image_bytes)
+    if preprocessed_pil is None: return "Error during preprocessing", "", ""
+    
+    # 2. ROUTE (Using preprocessed image)
+    mode, ar, bc = original_classify_input(preprocessed_pil)
+    status = f"Mode: {mode.upper()} (AR: {ar:.2f}, Blobs: {bc})"
     
-    # 1. Routing
-    if manual_mode == "Automatic":
-        mode, ar, bc = classify_input_type(image)
-        routing_status = f"Auto: {mode.upper()} (AR: {ar:.2f}, Blobs: {bc})"
-    else:
-        mode = manual_mode.lower()
-        routing_status = f"Manual: {mode.upper()}"
-        
     try:
         if mode == "character" and cnn_engine.available:
-            result = cnn_engine.predict(image)
-            return result["text"], routing_status, "CNN Classifier", ""
+            result = cnn_engine.predict(preprocessed_pil)
+            return result["text"], status, "CNN Classifier"
         else:
-            # Word Recognition
-            pixel_values = processor(image.convert("RGB"), return_tensors="pt").pixel_values.to(device)
+            pixel_values = processor(preprocessed_pil, return_tensors="pt").pixel_values.to(device)
             with torch.no_grad():
-                outputs = model.generate(
-                    pixel_values,
-                    num_beams=4,
-                    max_length=128,
-                    early_stopping=True
-                )
+                outputs = model.generate(pixel_values, num_beams=4, max_length=128, early_stopping=True)
             text = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-            return text, routing_status, "TrOCR + LoRA", ""
+            return text, status, "TrOCR + LoRA"
     except Exception as e:
-        return f"Error: {str(e)}", "Inference Failed", "None", ""
+        return f"Error: {str(e)}", "Failed", "None"
 
 # --- UI ---
-CSS = """
-.gradio-container { background: #0f172a; color: white; font-family: 'Inter', sans-serif; }
-.panel { background: rgba(30, 41, 59, 0.8); border-radius: 20px; padding: 20px; border: 1px solid #334155; }
-.result-box { font-size: 2.5rem !important; font-weight: bold; color: #818cf8; text-align: center; }
-"""
-
-with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
-    gr.Markdown("# 🕉️ DevGen OCR — Professional Engine")
-    with gr.Row(elem_classes="panel"):
+with gr.Blocks(css=".gradio-container {background: #0f172a; color: white;}") as demo:
+    gr.Markdown("# 🕉️ DevGen OCR — Full Fidelity Suite")
+    with gr.Row():
         with gr.Column():
-            input_img = gr.Image(type="pil", label="Input Handwriting")
-            mode_sel = gr.Radio(["Automatic", "Word", "Character"], value="Automatic", label="Mode")
-            run_btn = gr.Button("Recognize", variant="primary")
+            inp = gr.Image(type="pil", label="Input Image")
+            btn = gr.Button("Recognize", variant="primary")
         with gr.Column():
-            output_text = gr.Textbox(label="Recognition Result", elem_classes="result-box")
-            status_lbl = gr.Markdown("Engine ready.")
-            engine_lbl = gr.Textbox(label="Model Used", interactive=False)
+            out = gr.Textbox(label="Result", interactive=False)
+            lbl = gr.Label(label="Engine Status")
+            eng = gr.Textbox(label="Model", interactive=False)
 
-    run_btn.click(predict, [input_img, mode_sel], [output_text, status_lbl, engine_lbl])
+    btn.click(predict, [inp], [out, lbl, eng])
 
 if __name__ == "__main__":
     demo.launch()

preprocessing.py

@@ -0,0 +1,53 @@
+import cv2
+import numpy as np
+from PIL import Image
+import io
+
+def bytes_to_cv2(image_bytes: bytes) -> np.ndarray:
+    nparr = np.frombuffer(image_bytes, np.uint8)
+    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    return img
+
+def cv2_to_pil(img: np.ndarray) -> Image.Image:
+    rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(rgb)
+
+def crop_to_foreground(img: np.ndarray, padding_ratio: float = 0.18) -> np.ndarray:
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    _, mask = cv2.threshold(blurred, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
+    kernel = np.ones((3, 3), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    mask = cv2.dilate(mask, kernel, iterations=1)
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if not contours: return img
+    h, w = gray.shape[:2]
+    min_area = max(12, int(h * w * 0.0001))
+    boxes = [cv2.boundingRect(contour) for contour in contours if cv2.contourArea(contour) >= min_area]
+    if not boxes: return img
+    x1, y1, x2, y2 = min(x for x,_,_,_ in boxes), min(y for _,y,_,_ in boxes), max(x+bw for x,_,bw,_ in boxes), max(y+bh for _,y,_,bh in boxes)
+    pad_x, pad_y = max(8, int((x2 - x1) * padding_ratio)), max(8, int((y2 - y1) * padding_ratio))
+    x1, y1, x2, y2 = max(0, x1 - pad_x), max(0, y1 - pad_y), min(w, x2 + pad_x), min(h, y2 + pad_y)
+    return img[y1:y2, x1:x2]
+
+def normalize_for_model(img: np.ndarray, target_height: int = 384, target_width: int = 384) -> np.ndarray:
+    h, w = img.shape[:2]
+    scale = min(target_height / h, target_width / w)
+    new_h, new_w = int(h * scale), int(w * scale)
+    resized = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * 255
+    y_offset, x_offset = (target_height - new_h) // 2, (target_width - new_w) // 2
+    canvas[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized
+    return canvas
+
+def preprocess_for_ocr(image_bytes: bytes) -> Image.Image:
+    img = bytes_to_cv2(image_bytes)
+    if img is None: return None
+    h, w = img.shape[:2]
+    aspect_ratio = w / float(h)
+    if aspect_ratio <= 1.55:
+        img = crop_to_foreground(img)
+    elif aspect_ratio > 2.2:
+        img = crop_to_foreground(img)
+        img = normalize_for_model(img)
+    return cv2_to_pil(img)