Update app.py

app.py

@@ -38,9 +38,10 @@ def visualize_prediction(pil_img, output_dict, threshold=0.7, id2label=None):
     boxes = output_dict["boxes"][keep].tolist()
     scores = output_dict["scores"][keep].tolist()
     labels = output_dict["labels"][keep].tolist()
+    print(labels)
     if id2label is not None:
         labels = [id2label[x] for x in labels]
-
+    res = dict(zip(labels, scores))
     plt.figure(figsize=(16, 10))
     plt.imshow(pil_img)
     ax = plt.gca()
@@ -49,7 +50,7 @@ def visualize_prediction(pil_img, output_dict, threshold=0.7, id2label=None):
         ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=color, linewidth=3))
         ax.text(xmin, ymin, f"{label}: {score:0.2f}", fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
     plt.axis("off")
-    return fig2img(plt.gcf())
+    return fig2img(plt.gcf()),res
 
 def detect_objects(model_name,image_input,threshold):
     print(type(image_input))
@@ -71,9 +72,9 @@ def detect_objects(model_name,image_input,threshold):
     processed_outputs = make_prediction(image, feature_extractor, model)
     
     #Visualize prediction
-    viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
+    viz_img,labels = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
     
-    return viz_img   
+    return viz_img,labels
         
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
@@ -116,11 +117,14 @@ with demo:
           with gr.Row():
             example_images = gr.Dataset(components=[img_input],
                                             samples=[["airport.jpg"],['football-match.jpg']])
-                
           img_but = gr.Button('Detect')
+
+        with gr.TabItem('Labels'):
+          with gr.Row():
+            label = gr.Label(label = 'Labels')
         
-    url_but.click(detect_objects,inputs=[options,url_input,slider_input],outputs=img_output_from_url,queue=True)
-    img_but.click(detect_objects,inputs=[options,img_input,slider_input],outputs=img_output_from_upload,queue=True)
+    url_but.click(detect_objects,inputs=[options,url_input,slider_input],outputs=[img_output_from_url,label],queue=True)
+    img_but.click(detect_objects,inputs=[options,img_input,slider_input],outputs=[img_output_from_upload,label],queue=True)
     example_images.click(fn=set_example_image,inputs=[example_images],outputs=[img_input])
     example_url.click(fn=set_example_url,inputs=[example_url],outputs=[url_input])
     
@@ -128,4 +132,4 @@ with demo:
     #gr.Markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=nickmuchi-object-detection-with-detr-and-yolos)")
 
     
-demo.launch(enable_queue=True)
+demo.launch(enable_queue=True,show_api=False)