Switching to mobilenet_v2

app.py

@@ -2,18 +2,20 @@
 # gradio demo                                                                                                     
 #                                                                                                                 
 import gradio as gr
-from transformers import ViTFeatureExtractor, ViTModel
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
-model = ViTModel.from_pretrained('akahana/vit-base-cats-vs-dogs')
+preprocessor = AutoImageProcessor.from_pretrained("google/mobilenet_v2_1.0_224")
+model = AutoModelForImageClassification.from_pretrained("google/mobilenet_v2_1.0_224")
 
 title = "Sandbox"
 description = "Place to try various models"
 
 def classify(image):
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    inputs = preprocessor(images=image, return_tensors="pt")
     outputs = model(**inputs)
-    return str(outputs)
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    return model.config.id2label[predicted_class_idx]
     
 demo = gr.Interface(fn=classify, inputs="image", outputs="text")