add label component

app.py

@@ -28,14 +28,13 @@ def classify_bird(img):
     with torch.no_grad():
         top_n_coordinates, concat_out, raw_logits, concat_logits, part_logits, top_n_index, top_n_prob = model(torch_images)
 
-        _, predict = torch.max(concat_logits, 1)
-        pred_id = predict.item()
-        bird_class = model.bird_classes[pred_id]
-        print(f"{bird_class=}")
+        probs = torch.softmax(concat_logits, 1)[0]
+        prob_dict = {bird_cls: float(prob) for bird_cls, prob in zip(model.bird_classes, probs)}
 
-    return bird_class
+    return prob_dict
 
 image_component = gr.Image(type="pil", label="Bird Image")
-demo = gr.Interface(fn=classify_bird, inputs=image_component, outputs="text")
+label_component = gr.Label(label="Classification result", num_top_classes=3)
+demo = gr.Interface(fn=classify_bird, inputs=image_component, outputs=label_component)
 
 demo.launch()