Update app.py

app.py

@@ -15,12 +15,13 @@ def segment(image):
     with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.nn.Softmax(dim=1)(logits)
-       labels = [(prob, model.config.id2label[idx]) for idx, prob in enumerate(probs[0])]
+       # labels = [(prob, model.config.id2label[idx]) for idx, prob in enumerate(probs[0])]
+       labels = {model.config.id2label[idx] : flaot(prob) for idx, prob in enumerate(probs[0])}
        print(labels)
     
     # model predicts one of the 1000 ImageNet classes
     predicted_label = logits.argmax(-1).item()
-    return model.config.id2label[predicted_label]
+    return labels # model.config.id2label[predicted_label]
     
 gr.Interface(fn=segment, inputs="image", outputs="text").launch()