Update app.py

app.py

@@ -1,35 +1,32 @@
 import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from PIL import Image
-import numpy as np
 import torch
+import numpy as np
+from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq
 
-# Load TrOCR model and processor
-processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
+# Load the model & processor
+model_name = "Murasajo/Llama-3.2-VL-Finetuned-on-HandwrittenText"
+processor = AutoProcessor.from_pretrained(model_name)
+model = AutoModelForVision2Seq.from_pretrained(model_name)
 
-# Function to extract text from handwritten images
+# Function to extract handwritten text
 def extract_text(image):
     try:
-        # Convert image to RGB if needed
+        # Convert input to PIL Image
         if isinstance(image, np.ndarray):  
-            if len(image.shape) == 2:  # If grayscale (H, W), convert to RGB
+            if len(image.shape) == 2:  # If grayscale (H, W), add channels
                 image = np.stack([image] * 3, axis=-1)
             image = Image.fromarray(image)
         else:
             image = Image.open(image).convert("RGB")
 
-        # Preprocessing (convert to grayscale for better OCR)
-        image = image.convert("L")  
-        image = image.resize((640, 640))  
-
-        # Process image
+        # Process image through model
         pixel_values = processor(images=image, return_tensors="pt").pixel_values
         generated_ids = model.generate(pixel_values)
         extracted_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
-        return extracted_text if extracted_text.strip() else "No text detected."
-    
+        return extracted_text
+
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -38,9 +35,9 @@ iface = gr.Interface(
     fn=extract_text,
     inputs="image",
     outputs="text",
-    title="Handwritten OCR Extractor",
-    description="Upload a handwritten image to extract text.",
+    title="Handwritten Text OCR",
+    description="Upload a handwritten document and extract text using AI.",
 )
 
-# Launch the app
+# Run the app
 iface.launch()