Update app.py

app.py

@@ -1,15 +1,15 @@
-# app.py
+# app.py (Enhanced to display both image and text)
 
 import gradio as gr
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from model import load_model
 from utils import preprocess_image, decode_predictions
 import os
 
 # Load the model (ensure the path is correct)
-MODEL_PATH = "./finetuned_recog_model.pth"
-FONT_PATH = "./NotoSansEthiopic-Regular.ttf"  # Update the path to your font
+MODEL_PATH = "finetuned_recog_model.pth"
+FONT_PATH = "NotoSansEthiopic-Regular.ttf"  # Update the path to your font
 
 # Check if model file exists
 if not os.path.exists(MODEL_PATH):
@@ -28,41 +28,36 @@ from matplotlib import font_manager as fm
 import matplotlib.pyplot as plt
 
 ethiopic_font = fm.FontProperties(fname=FONT_PATH, size=15)
+pil_font = ImageFont.truetype(FONT_PATH, size=20)
 
-def recognize_text(image: Image.Image):
+def recognize_and_overlay(image: Image.Image):
     """
-    Function to recognize text from an image.
+    Function to recognize text and overlay it on the image.
+    Returns both the modified image and the recognized text.
     """
-    # Preprocess the image
-    input_tensor = preprocess_image(image).unsqueeze(0).to(device)  # [1, 3, 224, 224]
+    recognized_text = recognize_text(image)
     
-    # Perform inference
-    with torch.no_grad():
-        log_probs = model(input_tensor)  # [H*W, 1, vocab_size]
+    # Create a copy of the image to draw on
+    image_with_text = image.copy()
     
-    # Decode predictions
-    recognized_texts = decode_predictions(log_probs)
+    # Overlay text on the image
+    draw = ImageDraw.Draw(image_with_text)
+    text_position = (10, 10)  # Top-left corner
+    text_color = (255, 0, 0)  # Red color
+    draw.text(text_position, f"Recognized: {recognized_text}", font=pil_font, fill=text_color)
     
-    return recognized_texts[0]
-
-def display_image_with_text(image: Image.Image, recognized_text: str):
-    """
-    Function to display the image with recognized text.
-    """
-    plt.figure(figsize=(6,6))
-    plt.imshow(image)
-    plt.axis('off')
-    plt.title(f"Recognized Text: {recognized_text}", fontproperties=ethiopic_font)
-    plt.show()
-    return plt
+    return image_with_text, recognized_text
 
 # Define Gradio Interface
 iface = gr.Interface(
-    fn=recognize_text,
-    inputs=gr.inputs.Image(type="pil"),
-    outputs=gr.outputs.Textbox(),
+    fn=recognize_and_overlay,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=[
+        gr.Image(type="pil", label="Image with Recognized Text"),
+        gr.Textbox(label="Recognized Text")
+    ],
     title="Amharic Text Recognition",
-    description="Upload an image containing Amharic text, and the model will recognize and display the text."
+    description="Upload an image containing Amharic text. The app will recognize and overlay the text on the image."
 )
 
 # Launch the Gradio app