update app.py

app.py

@@ -1,19 +1,34 @@
 import gradio as gr
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch
+from PIL import Image
+import torchvision.transforms as transforms
 
-model_name = "./model.pkl"
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = torch.load("model.pkl", map_location=torch.device("cpu"))
+model.eval()
 
 
-def classify_text(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    predictions = outputs.logits.argmax(dim=1)
+def preprocess_image(image):
+    transform = transforms.Compose(
+        [
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ]
+    )
+    image_tensor = transform(image)
+    return image_tensor.unsqueeze(0)
 
-    return f"Predicted class: {predictions.item()}"
 
+def classify_image(image):
+    input_tensor = preprocess_image(image)
 
-iface = gr.Interface(fn=classify_text, inputs="text", outputs="text")
+    with torch.no_grad():
+        outputs = model(input_tensor)
+        predictions = torch.argmax(outputs, dim=1).item()
+
+    return f"Predicted class: {predictions}"
+
+
+iface = gr.Interface(fn=classify_image, inputs="image", outputs="text")
 
 iface.launch()

requirements.txt

@@ -1,2 +1,3 @@
 gradio==2.1.2
-transformers==4.13.0
+torch 
+torchvision