debug

app.py

@@ -4,9 +4,11 @@ import gradio as gr
 import numpy as np
 import torch
 from transformers import CLIPProcessor, CLIPModel
+from PIL import Image
 
 IMAGENET_CLASSES_FILE = "imagenet-classes.txt"
 EXAMPLES = ["dog.jpeg", "car.png"]
+RESIZED_IMAGE_SIZE = 640
 
 MARKDOWN = """
 # Zero-Shot Image Classification with MetaCLIP
@@ -24,6 +26,18 @@ def load_text_lines(file_path: str) -> List[str]:
         return [line.rstrip() for line in lines]
 
 
+def resize_image(input_image):
+    aspect_ratio = input_image.width / input_image.height
+    if input_image.width > input_image.height:
+        new_width = RESIZED_IMAGE_SIZE
+        new_height = int(RESIZED_IMAGE_SIZE / aspect_ratio)
+    else:
+        new_height = RESIZED_IMAGE_SIZE
+        new_width = int(RESIZED_IMAGE_SIZE * aspect_ratio)
+
+    return input_image.resize((new_width, new_height), Image.LANCZOS)
+
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = CLIPModel.from_pretrained("facebook/metaclip-b32-400m").to(device)
 processor = CLIPProcessor.from_pretrained("facebook/metaclip-b32-400m")
@@ -33,7 +47,7 @@ imagenet_classes = load_text_lines(IMAGENET_CLASSES_FILE)
 def classify_image(input_image) -> str:
     inputs = processor(
         text=imagenet_classes,
-        images=input_image,
+        images=resize_image(input_image),
         return_tensors="pt",
         padding=True).to(device)
     outputs = model(**inputs)