Update app.py

app.py

@@ -1,23 +1,33 @@
 import gradio
-from transformers import AutoImageProcessor, MobileNetV2Model
 import torch
-from datasets import load_dataset
+from transformers import AutoImageProcessor, MobileNetV2ForImageClassification
 
 image_processor = AutoImageProcessor.from_pretrained("Aruno/gemini-beauty")
-model = MobileNetV2Model.from_pretrained("Aruno/gemini-beauty")
+model = MobileNetV2ForImageClassification.from_pretrained("Aruno/gemini-beauty")
+
 
 def inference(img):
     inputs = image_processor(img, return_tensors="pt")
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = model(**inputs).logits[0]
+        outputs = outputs.softmax(dim=0)
+        outputs = {
+            "attractive": outputs[0],
+            "normal": outputs[1],
+            "ugly": outputs[2],
+            "very attractive": outputs[3],
+            "very_ugly": outputs[4],
+        }
     return outputs
 
+
 iface = gradio.Interface(
-  fn=inference,
-  inputs='image',
-  outputs='image',
-  title='Your Attractivness', 
-  description='Check your attractivness',
-  examples=["face_01.jpg", "face_02.jpg"])  
+    fn=inference,
+    inputs="image",
+    outputs="label",
+    title="Your Attractivness",
+    description="Check your attractivness",
+    examples=["face_01.jpg", "face_02.jpg"],
+)
 
-iface.launch()
+iface.launch()