Update README.md

README.md

@@ -1,23 +1,33 @@
 ```python
-from PIL import Image
+import torch
+import numpy as np
+from PIL import Image, ImageOps
 import matplotlib.pyplot as plt
+from torchvision import transforms as T
 from transformers import VisionEncoderDecoderModel, AutoTokenizer
 
 tokenizer = AutoTokenizer.from_pretrained('tirthadagr8/CustomOCR')
 model=VisionEncoderDecoderModel.from_pretrained('tirthadagr8/CustomOCR')
 import torch
 from torchvision import transforms as T
-simple_transforms=T.Compose([
-            T.Resize((224,224)),
-            T.ToTensor(),
-            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
-        ])
+def resize_with_padding(image, target_size=(224, 224)):
+    # Resize to fit within target_size while preserving aspect ratio
+    image.thumbnail((target_size[0], target_size[1]))
+    delta_w = target_size[0] - image.width
+    delta_h = target_size[1] - image.height
+    padding = (delta_w//2, delta_h//2, delta_w - (delta_w//2), delta_h - (delta_h//2))
+    padded_img = ImageOps.expand(image, padding, fill="white")
+    transform = T.Compose([
+        T.ToTensor(),  # Convert to tensor and scale to [0, 1]
+        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # Normalize
+    ])
+    return transform((padded_img))
 
-path="image.jpg"
-img=simple_transforms(Image.open(path))
+path="0106.jpg"
+img=resize_with_padding(Image.open(path))
 model.eval()
 with torch.no_grad():
     print(tokenizer.batch_decode(model.cuda().generate(img.unsqueeze(0).cuda()),skip_special_tokens=True))
 
-plt.imshow(Image.open(path).resize((224,224)))
+plt.imshow(img.permute(1,2,0).detach().cpu().numpy())
 ```