Update app.py

app.py

@@ -16,14 +16,15 @@ def segment(image):
        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 = {model.config.id2label[idx] : flaot(prob) for idx, prob in enumerate(probs[0])}
+       labels = {model.config.id2label[idx] : float(prob) for idx, prob in enumerate(probs[0])}
        print(labels)
     
     # model predicts one of the 1000 ImageNet classes
-    predicted_label = logits.argmax(-1).item()
+    # predicted_label = logits.argmax(-1).item()
     return labels # model.config.id2label[predicted_label]
     
-gr.Interface(fn=segment, inputs="image", outputs="text").launch()
+gr.Interface(fn=segment, inputs="image", outputs="label").launch()
+#gr.Interface(fn=segment, inputs="image", outputs="text").launch()
     
 #    with torch.no_grad():
 #        prediction = torch.nn.functional.softmax(model(**inputs)[0], dim=0)