doleron
diff --git a/‎misc/kids_detection.png
1.32 MB b/‎misc/kids_detection.png
1.32 MB
diff --git a/‎misc/kids_detection_output.png
398 KB b/‎misc/kids_detection_output.png
398 KB
diff --git a/‎python/yolo-tiny.py
Lines changed: 35 additions & 8 deletions b/‎python/yolo-tiny.py
Lines changed: 35 additions & 8 deletions
@@ -7,8 +7,17 @@
 
 # step 2 - feed a 640x640 image to get predictions
 
-image = cv2.imread('misc/car.jpg')
-blob = cv2.dnn.blobFromImage(image, 1/255.0, (640, 640), swapRB=True)
+def format_yolov5(frame):
+
+    row, col, _ = frame.shape
+    _max = max(col, row)
+    result = np.zeros((_max, _max, 3), np.uint8)
+    result[0:row, 0:col] = frame
+    return result
+
+image = cv2.imread('misc/corrida-crianças-1280x720.jpg')
+input_image = format_yolov5(image) # making the image square
+blob = cv2.dnn.blobFromImage(input_image , 1/255.0, (640, 640), swapRB=True)
 net.setInput(blob)
 predictions = net.forward()
 
@@ -20,6 +29,10 @@
 
 output_data = predictions[0]
 
+image_width, image_height, _ = input_image.shape
+x_factor = image_width / 640
+y_factor =  image_height / 640
+
 for r in range(25200):
     row = output_data[r]
     confidence = row[4]
@@ -35,23 +48,37 @@
             class_ids.append(class_id)
 
             x, y, w, h = row[0].item(), row[1].item(), row[2].item(), row[3].item() 
-            left = int(x - 0.5 * w)
-            top = int(y - 0.5 * h)
-            box = np.array([left, top, int(w), int(h)])
+            left = int((x - 0.5 * w) * x_factor)
+            top = int((y - 0.5 * h) * y_factor)
+            width = int(w * x_factor)
+            height = int(h * y_factor)
+            box = np.array([left, top, width, height])
             boxes.append(box)
 
 class_list = []
 with open("config_files/classes.txt", "r") as f:
     class_list = [cname.strip() for cname in f.readlines()]
 
-for i in range(len(class_ids)):
+indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.25, 0.45) 
+
+result_class_ids = []
+result_confidences = []
+result_boxes = []
+
+for i in indexes:
+    result_confidences.append(confidences[i])
+    result_class_ids.append(class_ids[i])
+    result_boxes.append(boxes[i])
+
+for i in range(len(result_class_ids)):
 
-    box = boxes[i]
-    class_id = class_ids[i]
+    box = result_boxes[i]
+    class_id = result_class_ids[i]
 
     cv2.rectangle(image, box, (0, 255, 255), 2)
     cv2.rectangle(image, (box[0], box[1] - 20), (box[0] + box[2], box[1]), (0, 255, 255), -1)
     cv2.putText(image, class_list[class_id], (box[0], box[1] - 10), cv2.FONT_HERSHEY_SIMPLEX, .5, (0,0,0))
 
+cv2.imwrite("misc/kids_detection.png", image)
 cv2.imshow("output", image)
 cv2.waitKey()