Update app.py

app.py

@@ -14,11 +14,13 @@ def segment(image):
     
     with torch.no_grad():
        logits = model(**inputs).logits
-       print(logits)
+       probs = torch.nn.Softmax(dim=1)(logits)
+       print(probs)
     
     # model predicts one of the 1000 ImageNet classes
     predicted_label = logits.argmax(-1).item()
     return model.config.id2label[predicted_label]
+    
 gr.Interface(fn=segment, inputs="image", outputs="text").launch()
     
 #    with torch.no_grad():