Custom Model

app.py

@@ -1,39 +1,33 @@
-import gradio as gr
-import os
 import torch
 from PIL import Image
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from huggingface_hub import hf_hub_download
+import os
 
-# Set up device
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# 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")
 
-# Load the fine-tuned model
-checkpoint_path = './checkpoint-2070'  # Path to your fine-tuned model checkpoint
-model = VisionEncoderDecoderModel.from_pretrained(checkpoint_path).to(device)
+# Set up the device (GPU or CPU)
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# Use the original model's processor (tokenizer and feature extractor)
+# Load the fine-tuned model and processor from the downloaded folder
+model = VisionEncoderDecoderModel.from_pretrained(checkpoint_folder).to(device)
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
 
 def ocr_image(image):
     """
-    Perform OCR on a single image.
-    :param image: PIL Image object.
-    :return: Extracted text from the image.
+    Perform OCR on an image using the loaded model.
+    :param image: Input PIL image.
+    :return: Extracted text.
     """
-    pixel_values = processor(image, return_tensors='pt').pixel_values.to(device)
+    # Preprocess image and generate OCR text
+    pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
     generated_ids = model.generate(pixel_values)
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return generated_text
 
-# Define the Gradio interface
-interface = gr.Interface(
-    fn=ocr_image,                      # Function to call for prediction
-    inputs=gr.inputs.Image(type="pil"), # Accept an image as input
-    outputs="text",                     # Return extracted text
-    title="OCR with TrOCR",
-    description="Upload an image, and the fine-tuned TrOCR model will extract the text for you."
-)
-
-# Launch the Gradio app
-if __name__ == "__main__":
-    interface.launch()
+# 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)

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers
+gradio
+pillow