Update app.py

app.py

@@ -12,13 +12,17 @@ import gradio as gr
 def segment(image):
     inputs = feature_extractor(image, return_tensors="pt")
     
-    with torch.no_grad():
-        logits = model(**inputs).logits
+    # with torch.no_grad():
+    #    logits = model(**inputs).logits
     
     # model predicts one of the 1000 ImageNet classes
-    predicted_label = logits.argmax(-1).item()
-    # print(model.config.id2label[predicted_label])
+    # predicted_label = logits.argmax(-1).item()
+    # return model.config.id2label[predicted_label]
+    
+    with torch.no_grad():
+        prediction = torch.nn.functional.softmax(model(**inputs)[0], dim=0)
     
-    return model.config.id2label[predicted_label]
+    return {model.config.id2label[i]: float(prediction[i]) for i in range(3)}  
 
-gr.Interface(fn=segment, inputs="image", outputs="text").launch()
+# gr.Interface(fn=segment, inputs="image", outputs="text").launch()
+gr.Interface(fn=segment, inputs="image", outputs="label").launch()