feat: returning both label probs

app.py

@@ -35,9 +35,13 @@ def predict_fn(image):
             output = model(image)
         
         probabilities = torch.exp(output)
-        top_p, top_class = probabilities.topk(1, dim=1)
+        print(probabilities)
+        # focused (0) round probability
+        focused_prob = round(probabilities[0][0].item(), 2)
+        # distracted (1) round probability
+        distracted_prob = round(probabilities[0][1].item(), 2)
 
         # return dictionary whose keys are labels and values are confidences
-        return {str(top_class.item()): str(top_p.item())}
+        return {'focused': focused_prob, 'distracted': distracted_prob}
 
 gr.Interface(predict_fn, gr.inputs.Image(type="pil", label="Input Image"), outputs="label").launch()