Update app.py

app.py

@@ -2,31 +2,43 @@ import gradio as gr
 from transformers import TrOCRProcessor, VisionEncoderDecoderModel
 from PIL import Image
 import torch
+from torchvision import transforms
 import matplotlib.pyplot as plt
-import spaces
 
 # Load TrOCR model
 processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
 model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
 
-@spaces.GPU
+def preprocess_image(image):
+    # Convert image to RGB
+    image = image.convert("RGB")
+    
+    # Resize and normalize the image to [0, 1]
+    transform = transforms.Compose([
+        transforms.Resize((384, 384)),  # Resize to the expected input size
+        transforms.ToTensor(),          # Convert to tensor and scale to [0, 1]
+    ])
+    pixel_values = transform(image).unsqueeze(0)  # Add batch dimension
+    return pixel_values
+
+def visualize_image(pixel_values):
+    # Convert tensor to numpy array and permute dimensions for visualization
+    image = pixel_values.squeeze().permute(1, 2, 0).numpy()
+    plt.imshow(image)
+    plt.title("Preprocessed Image")
+    plt.show()
+
 def recognize_text(image):
     try:
-        # Convert image to RGB if it's not already
-        image = image.convert("RGB")
-        print("Image converted to RGB.")
-
         # Preprocess the image
-        pixel_values = processor(images=image, return_tensors="pt").pixel_values
+        pixel_values = preprocess_image(image)
         print("Image preprocessed. Pixel values shape:", pixel_values.shape)
 
         # Visualize preprocessed image
-        plt.imshow(pixel_values.squeeze().permute(1, 2, 0))
-        plt.title("Preprocessed Image")
-        plt.show()
+        visualize_image(pixel_values)
 
         # Generate text from the image
-        with torch.no_grad():  # Disable gradient calculation for inference
+        with torch.no_grad():
             generated_ids = model.generate(pixel_values)
             print("Generated IDs:", generated_ids)