Update app.py

app.py

@@ -11,20 +11,46 @@ yolov7_custom_weights = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
 
 model = torch.hub.load('Owaiskhan9654/yolov7-1:main',model='custom', path_or_model=yolov7_custom_weights, force_reload=True)  # Github repository https://github.com/Owaiskhan9654
 
-def object_detection(im, size=640):
-    results = model(im) 
+def object_detection(
+    image: gr.inputs.Image = None,
+    model_path: gr.inputs.Dropdown = None,
+    image_size: gr.inputs.Slider = 640,
+    conf_threshold: gr.inputs.Slider = 0.25,
+    iou_threshold: gr.inputs.Slider = 0.45,):
+  
+
+    results = model(image) 
+
     results.render()  
-    return Image.fromarray(results.imgs[0])
+    count_dict = results.pandas().xyxy[0]['name'].value_counts().to_dict()
+
+
+    if len(count_dict)>0:
+      return Image.fromarray(results.imgs[0]),str(count_dict)
+    else:
+        return Image.fromarray(results.imgs[0]),'No object Found. Add more Custom classes in the training set'
+
 
 title = "Yolov7 Custom"
 
-image = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Upload Image", optional=False)
+# image = gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Upload Image", optional=False)
+
+inputs = [
+    gr.inputs.Image(shape=(640, 640), image_mode="RGB", source="upload", label="Upload Image", optional=False),
+    gr.inputs.Dropdown(["best.pt",], 
+                       default="best.pt", label="Model"),
+    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
+]
 outputs = gr.outputs.Image(type="pil", label="Output Image")
+outputs_cls = gr.Label(label= "Categories Detected Proportion Statistics" )
+
 
 Custom_description="<center>Custom Training Performed on Kaggle <a href='https://www.kaggle.com/code/owaiskhan9654/shelf-object-detection-yolov7-pytorch/notebook' style='text-decoration: underline' target='_blank'>Link</a> </center><br> <center>Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors </center> <br> on around 140 general items in Stores"
 
 Footer = (
-    "<center>Model Trained by: Owais Ahmad Data Scientist at <b> Thoucentric </b> <a href=\"https://www.linkedin.com/in/owaiskhan9654/\">Visit Profile</a> <br></center>"
+    "<br><br><br><br><center>Model Trained by: Owais Ahmad Data Scientist at <b> Thoucentric </b> <a href=\"https://www.linkedin.com/in/owaiskhan9654/\">Visit Profile</a> <br></center>"
     
     "<center> Model Trained Kaggle Kernel <a href=\"https://www.kaggle.com/code/owaiskhan9654/shelf-object-detection-yolov7-pytorch/notebook\">Link</a> <br></center>"
         
@@ -34,15 +60,17 @@ Footer = (
 
 examples1=[["Image1.jpg"],["Image2.jpg"],["Image3.jpg"],["Image4.jpg"],["Image5.jpg"],["Image6.jpg"]]
 
-Top_Title="<center>Yolov7 🚀 Custom Trained by <a href='https://www.linkedin.com/in/owaiskhan9654/' style='text-decoration: underline' target='_blank'>Owais Ahmad </center></a> on around 140 general items in Stores"
+Top_Title="<br><br><br><center>Yolov7 🚀 Custom Trained by <a href='https://www.linkedin.com/in/owaiskhan9654/' style='text-decoration: underline' target='_blank'>Owais Ahmad </center></a> on around 140 general items in Stores"
 css = ".output-image, .input-image {height: 50rem !important; width: 100% !important;}"
 css = ".image-preview {height: auto !important;}"
 
 gr.Interface(
     fn=object_detection,
-    inputs=image,
-    outputs=outputs,
+    inputs=inputs,
+    outputs=[outputs,outputs_cls],
     title=Top_Title,
     description=Custom_description,
     article=Footer,
+    cache= True,
+    allow_flagging=False,
     examples=examples1).launch()