Update app.py

app.py

@@ -1,67 +1,87 @@
 import gradio as gr
-from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from PIL import Image
 import torch
-import traceback
+from PIL import Image
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+import logging
 
-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
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
-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
+# Initialize TrOCR model and processor
+try:
+    processor = TrOCRProcessor.from_pretrained('microsoft/trocr-base-handwritten')
+    model = VisionEncoderDecoderModel.from_pretrained('microsoft/trocr-base-handwritten')
+    if torch.cuda.is_available():
+        model.to('cuda')
+except Exception as e:
+    logger.error(f"Error loading model: {e}")
+    raise
 
-def extract_text_from_image(image):
-    """Extract text from an uploaded image using Hugging Face TrOCR model."""
+def process_image(image):
+    """Process image and extract text using TrOCR"""
     try:
-        if image is None:
-            return "Error: No image provided"
-            
-        # Load model and processor
-        processor, model = load_model()
-        
-        # Preprocess image
-        image = preprocess_image(image)
+        # Convert to RGB if needed
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
         
-        # Extract text
+        # Prepare image for model
         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]
+            pixel_values = pixel_values.to('cuda')
         
-        return extracted_text.strip()
+        # Generate text
+        generated_ids = model.generate(pixel_values, max_length=128)
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
         
+        return generated_text.strip()
+    except Exception as e:
+        logger.error(f"Error processing image: {e}")
+        return f"Error processing image: {str(e)}"
+
+def analyze_image(input_image):
+    """Main function to handle image analysis"""
+    if input_image is None:
+        return "Please upload an image."
+    
+    try:
+        # Open and process image
+        image = Image.open(input_image)
+        
+        # Extract text
+        extracted_text = process_image(image)
+        
+        # Format response
+        response = f"""📝 Extracted Text:
+{'-' * 40}
+{extracted_text}
+{'-' * 40}
+
+📊 Statistics:
+• Characters: {len(extracted_text)}
+• Words: {len(extracted_text.split())}
+"""
+        return response
     except Exception as e:
-        error_msg = f"Error processing image: {str(e)}\n{traceback.format_exc()}"
-        return error_msg
+        logger.error(f"Error in analysis: {e}")
+        return f"Error analyzing image: {str(e)}"
 
-# Create Gradio Interface
-interface = gr.Interface(
-    fn=extract_text_from_image,
-    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.",
-    examples=["sample1.jpg", "sample2.jpg"]  # Add example images if you have them
+# Create Gradio interface
+demo = gr.Interface(
+    fn=analyze_image,
+    inputs=gr.Image(type="filepath", label="Upload Image"),
+    outputs=gr.Textbox(label="Extracted Text", lines=10),
+    title="📷 Smart OCR Text Extractor",
+    description="""
+    Extract text from images using Microsoft's TrOCR model.
+    Supports handwritten and printed text.
+    """,
+    theme=gr.themes.Soft(),
+    examples=[
+        ["example1.jpg"],
+        ["example2.png"]
+    ]
 )
 
 if __name__ == "__main__":
-    interface.launch(share=True)
+    demo.launch()