Update app.py

app.py

@@ -52,6 +52,12 @@ def visualize_prediction(pil_img, output_dict, threshold=0.5, id2label=None):
         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())
+    
+def get_original_image(url_input):
+    if validators.url(url_input):
+        image = Image.open(requests.get(url_input, stream=True).raw)
+        
+        return image
 
 def detect_objects(model_name,url_input,image_input,webcam_input,threshold):
     
@@ -60,14 +66,14 @@ def detect_objects(model_name,url_input,image_input,webcam_input,threshold):
     
     model = YolosForObjectDetection.from_pretrained(model_name)
     
+    
     if validators.url(url_input):
-        image = Image.open(requests.get(url_input, stream=True).raw)
+        image = get_original_image(url_input)
         #Make prediction
         processed_outputs = make_prediction(image, feature_extractor, model)
         #Visualize prediction
         viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
-        
-        return viz_img, image 
+         
         
     elif image_input:
         image = image_input
@@ -81,7 +87,7 @@ def detect_objects(model_name,url_input,image_input,webcam_input,threshold):
     #Visualize prediction
     viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
     
-    return viz_img, None   
+    return viz_img
         
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
@@ -136,7 +142,7 @@ with demo:
                     img_output_from_url = gr.Image(shape=(750,750))
                 
             with gr.Row():
-                example_url = gr.Dataset(components=[url_input],samples=[[str(url)] for url in urls])
+                example_url = gr.Dataset(components=[url_input,original_image],samples=[[str(url)] for url in urls],[])
             
             url_but = gr.Button('Detect')
      
@@ -163,7 +169,7 @@ with demo:
     img_but.click(detect_objects,inputs=[options,url_input,img_input,web_input,slider_input],outputs=[img_output_from_upload,original_image],queue=True)
     cam_but.click(detect_objects,inputs=[options,url_input,img_input,web_input,slider_input],outputs=[img_output_from_webcam,original_image],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])
+    example_url.click(fn=[set_example_url,get_original_image],inputs=[example_url,example_url],outputs=[url_input,original_image])
     
 
     gr.Markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=nickmuchi-face-mask-detections-with-yolos)")