Update app.py

app.py

@@ -3,33 +3,35 @@ import cv2
 import numpy as np
 import gradio as gr
 
-# load model
+
 model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
 
-# set model parameters
-model.conf = 0.25  # NMS confidence threshold
-model.iou = 0.45  # NMS IoU threshold
-model.agnostic = False  # NMS class-agnostic
-model.multi_label = False  # NMS multiple labels per box
-model.max_det = 1000  # maximum number of detections per image
+
+model.conf = 0.25  
+model.iou = 0.45  
+model.agnostic = False  
+model.multi_label = False  
+model.max_det = 1000
 
 
 def detect(img):
 
-# perform inference
+
     results = model(img, size=640)
 
-# inference with test time augmentation
-    results = model(img, augment=True)
-# parse results
     predictions = results.pred[0]
     boxes = predictions[:, :4] # x1, y1, x2, y2
     scores = predictions[:, 4]
     categories = predictions[:, 5]
+    dfResults = results.pandas().xyxy[0]
+    return drawRectangles(image, dfResults[['xmin', 'ymin', 'xmax','ymax']].astype(int))
+
+def drawRectangles(image, dfResults):
+    for index, row in dfResults.iterrows():
+      print( (row['xmin'], row['ymin']))
+      image = cv2.rectangle(image, (row['xmin'], row['ymin']), (row['xmax'], row['ymax']), (255, 0, 0), 2)
+    return image    
 
-    # save image as numpy array
-    return results.numpy()
-# show detection bounding boxes on image
 
 img = gr.inputs.Image(shape=(192, 192))