Update app.py

app.py

@@ -1,6 +1,9 @@
 import torch
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
+import cv2
+import numpy as np
+from craft_text_detector import Craft
 import gradio as gr
 
 # Force CPU usage
@@ -9,17 +12,28 @@ torch.set_default_device('cpu')
 # Load model and processor
 processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
 model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
+craft = Craft(output_dir=None, crop_type="box", cuda=False)
 
 def recognize_handwritten(image):
-    # Convert uploaded image to RGB
-    image = image.convert("RGB")
-    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    # Convert Gradio image to OpenCV format
+    image_np = np.array(image)
+    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
     
-    # Generate text
-    generated_ids = model.generate(pixel_values)
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    # Detect text regions with Craft
+    result = craft.detect_text(image=image_cv)
+    boxes = result["boxes"]
+    pil_image = Image.fromarray(cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB))
+    texts = []
     
-    return f"Recognized text: {generated_text}"
+    for box in boxes:
+        crop = pil_image.crop([box[0][0], box[0][1], box[2][0], box[2][1]])
+        pixel_values = processor(images=crop, return_tensors="pt").pixel_values
+        generated_ids = model.generate(pixel_values)
+        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        texts.append(text)
+    
+    text_data = " ".join(texts) if texts else "No text detected"
+    return f"Recognized text: {text_data}"
 
 # Create Gradio interface
 interface = gr.Interface(
@@ -31,4 +45,7 @@ interface = gr.Interface(
 )
 
 # Launch the app
-interface.launch()
+interface.launch()
+
+# Cleanup
+craft.unload_craftnet_model()