Update app.py

app.py

@@ -1,51 +1,54 @@
-import torch
+from hezar.models import Model
+from hezar.utils import load_image, draw_boxes
 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
+import numpy as np
+from PIL import Image
 
-# Force CPU usage
-torch.set_default_device('cpu')
-
-# Load model and processor
+# Load models on CPU (Hugging Face Spaces default)
+craft_model = Model.load("hezarai/CRAFT", device="cpu")
 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)
+trocr_model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
 
-def recognize_handwritten(image):
-    # Convert Gradio image to OpenCV format
+def recognize_handwritten_text(image):
+    # Convert Gradio image to format compatible with hezar
     image_np = np.array(image)
-    image_cv = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    processed_image = load_image(image_np)
+    
+    # Detect text regions with CRAFT
+    outputs = craft_model.predict(processed_image)
+    if not outputs or "boxes" not in outputs[0]:
+        return Image.fromarray(processed_image), "No text detected"
     
-    # 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))
+    boxes = outputs[0]["boxes"]
+    pil_image = Image.fromarray(processed_image)
     texts = []
     
+    # Recognize text in each detected region
     for box in boxes:
-        crop = pil_image.crop([box[0][0], box[0][1], box[2][0], box[2][1]])
+        x_min, y_min, x_max, y_max = box[0][0], box[0][1], box[2][0], box[2][1]
+        crop = pil_image.crop((x_min, y_min, x_max, y_max))
         pixel_values = processor(images=crop, return_tensors="pt").pixel_values
-        generated_ids = model.generate(pixel_values)
+        generated_ids = trocr_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}"
+    # Draw boxes on the image
+    result_image = draw_boxes(processed_image, boxes)
+    result_pil = Image.fromarray(result_image)
+    
+    # Join recognized texts
+    text_data = " ".join(texts) if texts else "No text recognized"
+    return result_pil, f"Recognized text: {text_data}"
 
 # Create Gradio interface
 interface = gr.Interface(
-    fn=recognize_handwritten,
+    fn=recognize_handwritten_text,
     inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Handwritten Text Recognition",
-    description="Upload an image containing handwritten text to recognize it."
+    outputs=[gr.Image(type="pil"), gr.Text()],
+    title="Handwritten Text Detection and Recognition",
+    description="Upload an image to detect and recognize handwritten text."
 )
 
 # Launch the app
-interface.launch()
-
-# Cleanup
-craft.unload_craftnet_model()
+interface.launch()