Update app.py

app.py

@@ -3,22 +3,60 @@ from transformers import VisionEncoderDecoderModel, TrOCRProcessor,AutoTokenizer
 from PIL import Image
 import torch
 
+def preprocess_image(image):
+    # Resize while maintaining aspect ratio
+    target_size = (224, 224)
+    original_size = image.size
+
+    # Calculate the new size while maintaining aspect ratio
+    aspect_ratio = original_size[0] / original_size[1]
+    if aspect_ratio > 1:  # Width is greater than height
+        new_width = target_size[0]
+        new_height = int(target_size[0] / aspect_ratio)
+    else:  # Height is greater than width
+        new_height = target_size[1]
+        new_width = int(target_size[1] * aspect_ratio)
+
+    # Resize the image
+    resized_img = image.resize((new_width, new_height))
+
+    # Calculate padding values
+    padding_width = target_size[0] - new_width
+    padding_height = target_size[1] - new_height
+
+    # Apply padding to center the resized image
+    pad_left = padding_width // 2
+    pad_top = padding_height // 2
+    pad_image = Image.new('RGB', target_size, (255, 255, 255))  # White background
+    pad_image.paste(resized_img, (pad_left, pad_top))
+    return pad_image
+
+
+# Load model directly
+from transformers import AutoTokenizer, AutoModel,ViTFeatureExtractor,TrOCRProcessor,VisionEncoderDecoderModel
 
 tokenizer = AutoTokenizer.from_pretrained("paudelanil/trocr-devanagari")
-model = VisionEncoderDecoderModel.from_pretrained("paudelanil/trocr-devanagari")
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+model1 = VisionEncoderDecoderModel.from_pretrained("paudelanil/trocr-devanagari")
+feature_extractor1 = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
+
+processor1 = TrOCRProcessor(feature_extractor=feature_extractor1, tokenizer=tokenizer)
+
+
+# tokenizer = AutoTokenizer.from_pretrained("paudelanil/trocr-devanagari")
+# model = VisionEncoderDecoderModel.from_pretrained("paudelanil/trocr-devanagari")
+# feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-model.to(device)
+model1.to(device)
 def predict(image):
     # Preprocess the image
     image = Image.open(image).convert("RGB")
     image = preprocess_image(image)
-    pixel_values = processor(image, return_tensors="pt").pixel_values.to(device)
+    pixel_values = processor1(image, return_tensors="pt").pixel_values.to(device)
     
     # Generate text from the image
-    generated_ids = model.generate(pixel_values)
-    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    generated_ids = model1.generate(pixel_values)
+    generated_text = processor1.batch_decode(generated_ids, skip_special_tokens=True)[0]
     
     return generated_text