Update app.py

app.py

@@ -12,17 +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()
-    # return model.config.id2label[predicted_label]
+    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():
-        prediction = torch.nn.functional.softmax(model(**inputs)[0], dim=0)
+#    with torch.no_grad():
+#        prediction = torch.nn.functional.softmax(model(**inputs)[0], dim=0)
     
-    return {model.config.id2label[i]: float(prediction[i]) for i in range(3)}  
+#    return {model.config.id2label[i]: float(prediction[i]) for i in range(3)}  
+#gr.Interface(fn=segment, inputs="image", outputs="label").launch()
 
-# gr.Interface(fn=segment, inputs="image", outputs="text").launch()
-gr.Interface(fn=segment, inputs="image", outputs="label").launch()