mejoras threshold

app.py

@@ -45,20 +45,34 @@ torch.hub.download_url_to_file('https://i.pinimg.com/originals/c2/ce/e0/c2cee056
 model = torch.hub.load('ultralytics/yolov5', 'custom', path='./best.pt')  # local model  o google colab
 #model = torch.hub.load('path/to/yolov5', 'custom', path='/content/yolov56.pt', source='local')  # local repo
 
-def yolo(im, size=640):
-    g = (size / max(im.size))  # gain
+
+def yolo(size, iou, conf, im):
+    '''Wrapper fn for gradio'''
+    g = (int(size) / max(im.size))  # gain
     im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS)  # resize
 
-    results = model(im)  # inference
-    results.render()  # updates results.imgs with boxes and labels
-    return Image.fromarray(results.imgs[0])
+    model.iou = iou
+    
+    model.conf = conf
+
+   
+    results2 = model(im)  # inference
+   
+    results2.render()  # updates results.imgs with boxes and labels
+    return Image.fromarray(results2.imgs[0])
+
+#------------ Interface-------------
 
 
-inputs = gr.inputs.Image(type='pil', label=" Imagen Original")
-outputs = gr.outputs.Image(type="pil", label="Resultado")
 
+in1 = gr.inputs.Radio(['640', '1280'], label="Tamaño de la imagen", default='640', type='value')
+in2 = gr.inputs.Slider(minimum=0, maximum=1, step=0.05, default=0.45, label='NMS IoU threshold')
+in3 = gr.inputs.Slider(minimum=0, maximum=1, step=0.05, default=0.50, label='Umbral o  threshold')
+in4 = gr.inputs.Image(type='pil', label="Original Image")
+
+out2 = gr.outputs.Image(type="pil", label="YOLOv5")
+#-------------- Text-----
 title = 'Trampas Barceló'
-        
 description = """
 <p>
 <center>
@@ -67,16 +81,15 @@ Sistemas de Desarrollado por Subsecretaría de Innovación del Municipio de Vice
 </center>
 </p>
 """
-              
-article = "<p style='text-align: center'>YOLOv5 is a family of compound-scaled object detection models trained on the COCO dataset, and includes " \
-          "simple functionality for Test Time Augmentation (TTA), model ensembling, hyperparameter evolution, " \
-          "and export to ONNX, CoreML and TFLite. <a href='https://colab.research.google.com/drive/1fbeB71yD09WK2JG9P3Ladu9MEzQ2rQad?usp=sharing'>Source code</a> |" \
-          "<a href='https://colab.research.google.com/drive/1FxaL8DcYgvjPrWfWruSA5hvk3J81zLY9?usp=sharing'>Colab Deploy</a> | <a href='https://github.com/ultralytics/yolov5'>PyTorch Hub</a></p>"
+article = "<p style='text-align: center'><a href='https://docs.google.com/presentation/d/1T5CdcLSzgRe8cQpoi_sPB4U170551NGOrZNykcJD0xU/edit?usp=sharing' target='_blank'>Para mas info, clik para ir al white paper</a></p><p style='text-align: center'><a href='https://drive.google.com/drive/folders/1owACN3HGIMo4zm2GQ_jf-OhGNeBVRS7l?usp=sharing ' target='_blank'>Google Colab Demo</a></p></center></p>"
           
-examples = [['ejemplo1.jpg'], ['ejemplo2.jpg']]
-gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch(
+examples = [['640',0.45, 0.75,'ejemplo1.jpg'], ['640',0.45, 0.75,'ejemplo2.jpg']]
+
+iface = gr.Interface(yolo, inputs=[in1, in2, in3, in4], outputs=out2, title=title, description=description, article=article, examples=examples,theme="huggingface", analytics_enabled=False).launch(
     debug=True)
 
+iface.launch()
+
 """For YOLOv5 PyTorch Hub inference with **PIL**, **OpenCV**, **Numpy** or **PyTorch** inputs please see the full [YOLOv5 PyTorch Hub Tutorial](https://github.com/ultralytics/yolov5/issues/36).