Update app.py

app.py

@@ -8,8 +8,7 @@ from PIL import Image
 from transformers import AutoFeatureExtractor, YolosForObjectDetection, DetrForObjectDetection
 import os
 
-
-os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
 
 # colors for visualization
 COLORS = [
@@ -34,12 +33,11 @@ def fig2img(fig):
     buf.seek(0)
     pil_img = Image.open(buf)
     basewidth = 750
-    wpercent = (basewidth/float(pil_img.size[0]))
-    hsize = int((float(pil_img.size[1])*float(wpercent)))
-    img = pil_img.resize((basewidth,hsize), Image.Resampling.LANCZOS) 
+    wpercent = (basewidth / float(pil_img.size[0]))
+    hsize = int((float(pil_img.size[1]) * float(wpercent)))
+    img = pil_img.resize((basewidth, hsize), Image.Resampling.LANCZOS)
     return img
 
-
 def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     keep = output_dict["scores"] > threshold
     boxes = output_dict["boxes"][keep].tolist()
@@ -47,9 +45,7 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
     labels = output_dict["labels"][keep].tolist()
     
     if id2label is not None:
-        
         labels = [id2label[x] for x in labels]
-        
 
     plt.figure(figsize=(50, 50))
     plt.imshow(img)
@@ -61,47 +57,39 @@ def visualize_prediction(img, output_dict, threshold=0.5, id2label=None):
             ax.text(xmin, ymin, f"{label}: {score:0.2f}", fontsize=60, bbox=dict(facecolor="yellow", alpha=0.8))
     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):
-    
-    #Extract model and feature extractor
+def detect_objects(model_name, url_input, image_input, webcam_input, threshold):
+    # Extract model and feature extractor
     feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-    
     if "yolos" in model_name:
         model = YolosForObjectDetection.from_pretrained(model_name)
     elif "detr" in model_name:
         model = DetrForObjectDetection.from_pretrained(model_name)
-    
     if validators.url(url_input):
         image = get_original_image(url_input)
-         
-    elif image_input:
+    elif image_input is not None:
         image = image_input
-        
-    elif webcam_input:
+    elif webcam_input is not None:
         image = webcam_input
-    
-    #Make prediction
+    else:
+        return None
+    # Make prediction
     processed_outputs = make_prediction(image, feature_extractor, model)
-    
-    #Visualize prediction
+    # Visualize prediction
     viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
-    
     return viz_img
-        
+
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
 
 def set_example_url(example: list) -> dict:
     return gr.Textbox.update(value=example[0]), gr.Image.update(value=get_original_image(example[0]))
 
-
 title = """<h1 id="title">License Plate Detection with YOLOS</h1>"""
 
 description = """
@@ -113,8 +101,8 @@ Links to HuggingFace Models:
 - [hustlv/yolos-small](https://huggingface.co/hustlv/yolos-small)  
 """
 
-models = ["nickmuchi/yolos-small-finetuned-license-plate-detection","nickmuchi/detr-resnet50-license-plate-detection"]
-urls = ["https://drive.google.com/uc?id=1j9VZQ4NDS4gsubFf3m2qQoTMWLk552bQ","https://drive.google.com/uc?id=1p9wJIqRz3W50e2f_A0D8ftla8hoXz4T5"]
+models = ["nickmuchi/yolos-small-finetuned-license-plate-detection", "nickmuchi/detr-resnet50-license-plate-detection"]
+urls = ["https://drive.google.com/uc?id=1j9VZQ4NDS4gsubFf3m2qQoTMWLk552bQ", "https://drive.google.com/uc?id=1p9wJIqRz3W50e2f_A0D8ftla8hoXz4T5"]
 images = [[path.as_posix()] for path in sorted(pathlib.Path('images').rglob('*.j*g'))]
 
 twitter_link = """
@@ -132,48 +120,40 @@ with demo:
     gr.Markdown(title)
     gr.Markdown(description)
     gr.Markdown(twitter_link)
-    options = gr.Dropdown(choices=models,label='Object Detection Model',value=models[0],show_label=True)
-    slider_input = gr.Slider(minimum=0.2,maximum=1,value=0.5,step=0.1,label='Prediction Threshold')
+    options = gr.Dropdown(choices=models, label='Object Detection Model', value=models[0], show_label=True)
+    slider_input = gr.Slider(minimum=0.2, maximum=1, value=0.5, step=0.1, label='Prediction Threshold')
     
     with gr.Tabs():
         with gr.TabItem('Image URL'):
             with gr.Row():
                 with gr.Column():
-                    url_input = gr.Textbox(lines=2,label='Enter valid image URL here..')
-                    original_image = gr.Image(shape=(750,750))
+                    url_input = gr.Textbox(lines=2, label='Enter valid image URL here..')
+                    original_image = gr.Image()
                     url_input.change(get_original_image, url_input, original_image)
                 with gr.Column():
-                    img_output_from_url = gr.Image(shape=(750,750))
-                
+                    img_output_from_url = gr.Image()
             with gr.Row():
-                example_url = gr.Examples(examples=urls,inputs=[url_input])
-                
-            
+                example_url = gr.Examples(examples=urls, inputs=[url_input])
             url_but = gr.Button('Detect')
      
         with gr.TabItem('Image Upload'):
             with gr.Row():
-                img_input = gr.Image(type='pil',shape=(750,750))
-                img_output_from_upload= gr.Image(shape=(750,750))
-                
-            with gr.Row(): 
-                example_images = gr.Examples(examples=images,inputs=[img_input])                            
-                                                   
-                
+                img_input = gr.Image(type='pil')
+                img_output_from_upload = gr.Image()
+            with gr.Row():
+                example_images = gr.Examples(examples=images, inputs=[img_input])
             img_but = gr.Button('Detect')
             
         with gr.TabItem('WebCam'):
             with gr.Row():
-                web_input = gr.Image(source='webcam',type='pil',shape=(750,750),streaming=True)
-                img_output_from_webcam= gr.Image(shape=(750,750))
-
+                web_input = gr.Image(source='webcam', type='pil', streaming=True)
+                img_output_from_webcam = gr.Image()
             cam_but = gr.Button('Detect')
             
-    url_but.click(detect_objects,inputs=[options,url_input,img_input,web_input,slider_input],outputs=[img_output_from_url],queue=True)
-    img_but.click(detect_objects,inputs=[options,url_input,img_input,web_input,slider_input],outputs=[img_output_from_upload],queue=True)
-    cam_but.click(detect_objects,inputs=[options,url_input,img_input,web_input,slider_input],outputs=[img_output_from_webcam],queue=True)
+    url_but.click(detect_objects, inputs=[options, url_input, img_input, web_input, slider_input], outputs=[img_output_from_url], queue=True)
+    img_but.click(detect_objects, inputs=[options, url_input, img_input, web_input, slider_input], outputs=[img_output_from_upload], queue=True)
+    cam_but.click(detect_objects, inputs=[options, url_input, img_input, web_input, slider_input], outputs=[img_output_from_webcam], queue=True)
 
     gr.Markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=nickmuchi-license-plate-detection-with-yolos)")
 
-    
-demo.launch(debug=True,enable_queue=True)
+demo.launch(debug=True, enable_queue=True)