Update app.py

app.py

@@ -1,20 +1,57 @@
 import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
-import requests
-from io import BytesIO
+import torch
+import traceback
 
-# Load TrOCR model and processor
-processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
-model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+def load_model():
+    """Load the TrOCR model and processor."""
+    processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
+    model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")
+    if torch.cuda.is_available():
+        model = model.to("cuda")
+    return processor, model
+
+def preprocess_image(image):
+    """Preprocess the input image."""
+    # Convert to RGB if needed
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    
+    # Resize if image is too large
+    max_size = 1000
+    if max(image.size) > max_size:
+        ratio = max_size / max(image.size)
+        new_size = tuple(int(dim * ratio) for dim in image.size)
+        image = image.resize(new_size, Image.LANCZOS)
+    
+    return image
 
 def extract_text_from_image(image):
     """Extract text from an uploaded image using Hugging Face TrOCR model."""
-    image = image.convert("RGB")
-    pixel_values = processor(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
+    try:
+        if image is None:
+            return "Error: No image provided"
+            
+        # Load model and processor
+        processor, model = load_model()
+        
+        # Preprocess image
+        image = preprocess_image(image)
+        
+        # Extract text
+        pixel_values = processor(image, return_tensors="pt").pixel_values
+        if torch.cuda.is_available():
+            pixel_values = pixel_values.to("cuda")
+            
+        generated_ids = model.generate(pixel_values)
+        extracted_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        
+        return extracted_text.strip()
+        
+    except Exception as e:
+        error_msg = f"Error processing image: {str(e)}\n{traceback.format_exc()}"
+        return error_msg
 
 # Create Gradio Interface
 interface = gr.Interface(
@@ -22,8 +59,9 @@ interface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=gr.Textbox(label="Extracted Text"),
     title="OCR Text Extractor",
-    description="Upload an image to extract text using Hugging Face's TrOCR model."
+    description="Upload an image to extract text using Hugging Face's TrOCR model.",
+    examples=["sample1.jpg", "sample2.jpg"]  # Add example images if you have them
 )
 
 if __name__ == "__main__":
-    interface.launch(share=True)
+    interface.launch(share=True)