Update app.py

app.py

@@ -6,50 +6,48 @@ import numpy as np
 import cv2
 from paddleocr import TextDetection
 
-MODEL_HUB_ID = "imperiusrex/Handwritten_model" # <--- MAKE SURE THIS IS CORRECT
+# --- Constants ---
+MODEL_HUB_ID = "imperiusrex/Handwritten_model"
+
+# --- Device ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# --- Load Models Globally ---
+print("🔄 Loading models...")
 
 processor = TrOCRProcessor.from_pretrained(MODEL_HUB_ID)
 model = VisionEncoderDecoderModel.from_pretrained(MODEL_HUB_ID)
-
-# Move model to appropriate device (GPU if available, else CPU)
-model.eval()
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
+model.eval()
 
 ocr_det_model = TextDetection(model_name="PP-OCRv5_server_det")
 
-# --- Inference Function for Gradio ---
+print("✅ Models loaded successfully.")
+
+# --- Inference Function ---
 def recognize_handwritten_text(image_input):
     if image_input is None:
         return "Please upload an image."
 
-    # Convert Gradio image input (numpy array) to PIL Image
     image_pil = Image.fromarray(image_input).convert("RGB")
 
-    # Perform text detection with PaddleOCR
-    # PaddleOCR expects a file path or numpy array
     detection_results = ocr_det_model.predict(image_input, batch_size=1)
 
     detected_polys = []
     for res in detection_results:
-        polys = res['dt_polys']
+        polys = res.get('dt_polys', [])
         if polys is not None:
             detected_polys.extend(polys.tolist())
 
     cropped_images = []
     if detected_polys:
-        img_np = np.array(image_pil) # Convert PIL to NumPy for OpenCV
+        img_np = np.array(image_pil)
 
-        for i, box in enumerate(detected_polys):
+        for box in detected_polys:
             box = np.array(box, dtype=np.float32)
 
-            width_a = np.linalg.norm(box[0] - box[1])
-            width_b = np.linalg.norm(box[2] - box[3])
-            height_a = np.linalg.norm(box[0] - box[3])
-            height_b = np.linalg.norm(box[1] - box[2])
-
-            width = int(max(width_a, width_b))
-            height = int(max(height_a, height_b))
+            width = int(max(np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[2] - box[3])))
+            height = int(max(np.linalg.norm(box[0] - box[3]), np.linalg.norm(box[1] - box[2])))
 
             dst_rect = np.array([
                 [0, 0],
@@ -60,9 +58,9 @@ def recognize_handwritten_text(image_input):
 
             M = cv2.getPerspectiveTransform(box, dst_rect)
             warped = cv2.warpPerspective(img_np, M, (width, height))
-            cropped_images.append(Image.fromarray(warped).convert("RGB")) # Convert back to PIL
+            cropped_images.append(Image.fromarray(warped).convert("RGB"))
 
-        cropped_images.reverse() # Apply reverse if that was intended based on your original code
+        cropped_images.reverse()
 
     recognized_texts = []
     if cropped_images:
@@ -73,23 +71,25 @@ def recognize_handwritten_text(image_input):
                 generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
                 recognized_texts.append(generated_text)
     else:
-        # Fallback if no text detected by PaddleOCR - process the whole image
         pixel_values = processor(images=image_pil, return_tensors="pt").pixel_values.to(device)
         with torch.no_grad():
             generated_ids = model.generate(pixel_values, max_new_tokens=64)
             generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
-            recognized_texts.append("No specific text regions detected, processing full image: " + generated_text)
+            recognized_texts.append("No text boxes detected. Full image OCR:\n" + generated_text)
 
     return "\n".join(recognized_texts)
 
-# --- Gradio Interface Setup ---
-iface = gr.Interface(
-    fn=recognize_handwritten_text,
-    inputs=gr.Image(type="numpy", label="Upload Handwritten Image"),
-    outputs="text",
-    title="Handwritten Text Recognition with TrOCR and PaddleOCR",
-    description="Upload an image with handwritten text to get it recognized. Uses PaddleOCR for text detection and TrOCR for recognition."
-)
-
-
-iface.launch()
+# --- Gradio Interface ---
+def build_interface():
+    return gr.Interface(
+        fn=recognize_handwritten_text,
+        inputs=gr.Image(type="numpy", label="Upload Handwritten Image"),
+        outputs="text",
+        title="✍️ Handwritten Text Recognition",
+        description="📷 Upload a handwritten image. Uses PaddleOCR (detection) + TrOCR (recognition).",
+    )
+
+# --- Launch App ---
+if __name__ == "__main__":
+    iface = build_interface()
+    iface.launch()