Restore: 1:1 copy of original backend routing and preprocessing logic

app.py

@@ -5,7 +5,7 @@ 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
 
 # --- CONFIGURATION ---
@@ -16,16 +16,9 @@ CNN_MODEL_PATH = "devanagari-cnn-classifier.pt"
 IS_SPACE = "SPACE_ID" in os.environ
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# --- MODEL INITIALIZATION ---
-print(f"System: Initializing DevGen OCR (Env: {'Hugging Face Space' if IS_SPACE else 'Local'})")
-
-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)
-
+# --- INITIALIZATION ---
+print(f"System: Restoring Original DevGen Logic...")
+processor = TrOCRProcessor.from_pretrained(BASE_MODEL_ID)
 base_model = VisionEncoderDecoderModel.from_pretrained(BASE_MODEL_ID)
 model = PeftModel.from_pretrained(base_model, ADAPTER_ID)
 model.to(device)
@@ -33,99 +26,99 @@ model.eval()
 
 cnn_engine = CharacterClassifier(model_path=CNN_MODEL_PATH, device=device)
 
-# --- REFINED ROUTING LOGIC ---
-def classify_input(image):
-    """Robustly detect if the image is a single character or a word."""
+# --- ORIGINAL ROUTING LOGIC (1:1 from image_router.py) ---
+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
+        size += 1
+        stack.extend([(y + 1, x), (y - 1, x), (y, x + 1), (y, x - 1)])
+    return size
+
+def _count_blobs(binary, min_size=10):
+    h, w = binary.shape
+    visited = np.zeros_like(binary, dtype=bool)
+    count = 0
+    for y in range(h):
+        for x in range(w):
+            if binary[y, x] and not visited[y, x]:
+                size = _flood_fill(binary, visited, y, x, h, w)
+                if size >= min_size:
+                    count += 1
+    return count
+
+def original_classify_input(image):
     gray = image.convert("L")
     arr = np.array(gray)
-    
-    # Binarize to find the ink
     threshold = min(arr.mean() * 0.75, 200)
     binary = (arr < threshold).astype(np.uint8)
     
-    # Find bounding box of ink
-    coords = np.column_stack(np.where(binary > 0))
-    if len(coords) == 0:
-        return "word", 0.5, "empty"
-        
-    y0, x0 = coords.min(axis=0)
-    y1, x1 = coords.max(axis=0)
+    rows = np.any(binary, axis=1)
+    cols = np.any(binary, axis=0)
+    if not rows.any() or not cols.any():
+        return "character", "no_ink"
     
-    w = x1 - x0 + 1
-    h = y1 - y0 + 1
+    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))
     
-    # Logic: Words in Devanagari are wide. Characters are nearly square.
-    # threshold 1.5 is the sweet spot for Devanagari characters vs words.
-    if aspect_ratio > 1.5:
-        return "word", 0.9, f"Wide (AR: {aspect_ratio:.1f})"
-    return "character", 0.9, f"Square (AR: {aspect_ratio:.1f})"
-
-def smart_predict(image, manual_mode):
-    if image is None:
-        return None, "Upload an image.", "None"
+    # decision logic (exact copy)
+    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
     
-    try:
-        # Determine mode
-        if manual_mode == "Automatic":
-            mode, conf, reason = classify_input(image)
-        else:
-            mode = manual_mode.lower()
-            conf, reason = 1.0, "Manual override"
-            
-        status = f"Mode: {mode.upper()} | {reason}"
-        
-        if mode == "character" and cnn_engine.available:
-            result = cnn_engine.predict(image)
-            return result["text"], status, "CNN Classifier"
-        else:
-            # Word Recognition Pipeline (TrOCR)
-            image_rgb = image.convert("RGB")
-            pixel_values = processor(image_rgb, return_tensors="pt").pixel_values.to(device)
-            with torch.no_grad():
-                generated_ids = model.base_model.generate(
-                    pixel_values=pixel_values,
-                    num_beams=4,
-                    max_new_tokens=64,
-                    early_stopping=True,
-                    decoder_start_token_id=model.config.decoder_start_token_id
-                )
-                text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            return text, status, "TrOCR + LoRA"
-            
-    except Exception as e:
-        return f"Error: {str(e)}", "System failure", "Error"
+    return "character" if is_character else "word", f"AR: {aspect_ratio:.2f}, Blobs: {blob_count}"
 
-# --- PREMIUM UI ---
-CSS = """
-@import url('https://fonts.googleapis.com/css2?family=Outfit:wght@400;600&family=Inter:wght@400;500&display=swap');
-.gradio-container { background: #0f172a !important; color: white !important; font-family: 'Inter', sans-serif !important; }
-.glass-card { background: rgba(30, 41, 59, 0.7) !important; backdrop-filter: blur(10px); border: 1px solid rgba(255,255,255,0.1); border-radius: 20px; padding: 25px; box-shadow: 0 20px 40px rgba(0,0,0,0.3); }
-h1 { font-family: 'Outfit', sans-serif; font-size: 2.5rem; background: linear-gradient(90deg, #818cf8, #c084fc); -webkit-background-clip: text; -webkit-fill-color: transparent; margin-bottom: 0.5rem; }
-.output-box { font-size: 2rem !important; text-align: center; background: rgba(0,0,0,0.2) !important; border-radius: 12px; border: 1px solid #475569 !important; }
-"""
+# --- MAIN INFERENCE ---
+def predict(image, manual_mode):
+    if image is None: return None, "Upload image.", ""
+    
+    if manual_mode == "Automatic":
+        mode, reason = original_classify_input(image)
+    else:
+        mode, reason = manual_mode.lower(), "Manual"
+    
+    if mode == "character" and cnn_engine.available:
+        # Use exact CNN preprocessing from CharacterClassifier
+        result = cnn_engine.predict(image)
+        return result["text"], f"Original Logic: {mode.upper()} ({reason})", "CNN Classifier"
+    else:
+        # Standard TrOCR inference
+        pixel_values = processor(image.convert("RGB"), return_tensors="pt").pixel_values.to(device)
+        with torch.no_grad():
+            gen_ids = model.base_model.generate(
+                pixel_values=pixel_values,
+                num_beams=4,
+                max_new_tokens=64,
+                decoder_start_token_id=model.config.decoder_start_token_id
+            )
+            text = processor.batch_decode(gen_ids, skip_special_tokens=True)[0]
+        return text, f"Original Logic: {mode.upper()} ({reason})", "TrOCR + LoRA"
 
-with gr.Blocks(css=CSS, theme=gr.themes.Default()) as demo:
-    with gr.Column(elem_classes="glass-card"):
-        gr.Markdown("# 🕉️ DevGen Smart OCR")
-        gr.Markdown("Automatic character and word recognition suite.")
-        
-        with gr.Row():
-            with gr.Column():
-                input_img = gr.Image(type="pil", label="Input Image")
-                mode_selector = gr.Radio(["Automatic", "Word", "Character"], value="Automatic", label="Recognition Mode")
-                submit_btn = gr.Button("Recognize", variant="primary")
-            
-            with gr.Column():
-                output_text = gr.Textbox(label="Recognition Result", elem_classes="output-box")
-                status_msg = gr.Markdown("Ready.")
-                model_info = gr.Textbox(label="Model Engine", interactive=False)
+# --- UI ---
+with gr.Blocks(theme=gr.themes.Default(), css=".gradio-container {background: #0f172a; color: white;}") as demo:
+    gr.Markdown("# 🕉️ DevGen OCR (Original Logic)")
+    with gr.Row():
+        with gr.Column():
+            inp = gr.Image(type="pil", label="Input")
+            mode = gr.Radio(["Automatic", "Word", "Character"], value="Automatic", label="Mode")
+            btn = gr.Button("Recognize", variant="primary")
+        with gr.Column():
+            out = gr.Textbox(label="Result")
+            status = gr.Markdown("Ready.")
+            eng = gr.Textbox(label="Engine", interactive=False)
 
-    submit_btn.click(
-        fn=smart_predict,
-        inputs=[input_img, mode_selector],
-        outputs=[output_text, status_msg, model_info]
-    )
+    btn.click(predict, [inp, mode], [out, status, eng])
 
 if __name__ == "__main__":
     demo.launch()