app

app.py

@@ -1,16 +1,12 @@
 import torch
 from PIL import Image
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
-from huggingface_hub import hf_hub_download
-import os
-
-# Load the model checkpoint and tokenizer files from Hugging Face Model Hub
-# checkpoint_folder = hf_hub_download(repo_id="Heramb26/tr-ocr-custom-checkpoints", filename="checkpoint-2070")
+import gradio as gr
 
 # Set up the device (GPU or CPU)
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Load the fine-tuned model and processor from the downloaded folder
+# Load the fine-tuned model and processor from the Hugging Face repository
 model = VisionEncoderDecoderModel.from_pretrained("Heramb26/TC-OCR-Custom").to(device)
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
 
@@ -26,8 +22,12 @@ def ocr_image(image):
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return generated_text
 
-# Example usage
-image_path = "path/to/your/image.jpg"  # Update with the path to your image
-image = Image.open(image_path)  # Open the image file using PIL
-extracted_text = ocr_image(image)  # Perform OCR on the image
-print("Extracted Text:", extracted_text)
+# Create a Gradio interface
+interface = gr.Interface(fn=ocr_image,  # Function to be called when an image is uploaded
+                         inputs=gr.inputs.Image(type="pil"),  # Input is an image file
+                         outputs="text",  # Output is extracted text
+                         title="OCR Inference",  # Title of the app
+                         description="Upload an image with handwritten text to extract the text.")  # Description
+
+# Launch the Gradio app
+interface.launch()