Skip to content

Commit 751cae6

Browse files
committed
fix: return empty bboxes when no detection
1 parent ce7130a commit 751cae6

File tree

1 file changed

+5
-16
lines changed

1 file changed

+5
-16
lines changed

motrackers/detectors/yolo.py

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,9 @@
22
import cv2 as cv
33
from motrackers.detectors.detector import Detector
44
from motrackers.utils.misc import load_labelsjson
5-
6-
75
class YOLOv3(Detector):
86
"""
97
YOLOv3 Object Detector Module.
10-
118
Args:
129
weights_path (str): path to network weights file.
1310
configfile_path (str): path to network configuration file.
@@ -17,41 +14,32 @@ class YOLOv3(Detector):
1714
draw_bboxes (bool): If True, assign colors for drawing bounding boxes on the image.
1815
use_gpu (bool): If True, try to load the model on GPU.
1916
"""
20-
2117
def __init__(self, weights_path, configfile_path, labels_path, confidence_threshold=0.5, nms_threshold=0.2, draw_bboxes=True, use_gpu=False):
2218
self.net = cv.dnn.readNetFromDarknet(configfile_path, weights_path)
2319
object_names = load_labelsjson(labels_path)
24-
2520
layer_names = self.net.getLayerNames()
2621
if cv.__version__ >= '4.6.0':
2722
self.layer_names = [layer_names[i - 1] for i in self.net.getUnconnectedOutLayers()]
2823
else:
29-
self.layer_names = [layer_names[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
30-
24+
self.layer_names = [layer_names[i - 1] for i in self.net.getUnconnectedOutLayers()]
25+
3126
self.scale_factor = 1/255.0
3227
self.image_size = (416, 416)
33-
3428
self.net = cv.dnn.readNetFromDarknet(configfile_path, weights_path)
3529
if use_gpu:
3630
self.net.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
3731
self.net.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
38-
3932
super().__init__(object_names, confidence_threshold, nms_threshold, draw_bboxes)
40-
4133
def forward(self, image):
4234
blob = cv.dnn.blobFromImage(image, self.scale_factor, self.image_size, swapRB=True, crop=False)
4335
self.net.setInput(blob)
4436
detections = self.net.forward(self.layer_names) # detect objects using object detection model
4537
return detections
46-
4738
def detect(self, image):
4839
if self.width is None or self.height is None:
4940
(self.height, self.width) = image.shape[:2]
50-
5141
detections = self.forward(image)
52-
5342
bboxes, confidences, class_ids = [], [], []
54-
5543
for output in detections:
5644
for detect in output:
5745
scores = detect[5:]
@@ -63,8 +51,9 @@ def detect(self, image):
6351
bboxes.append([x, y, w, h])
6452
confidences.append(float(confidence))
6553
class_ids.append(class_id)
66-
54+
if(bboxes == []):
55+
return [], [], []
6756
indices = cv.dnn.NMSBoxes(bboxes, confidences, self.confidence_threshold, self.nms_threshold).flatten()
6857
class_ids = np.array(class_ids).astype('int')
6958
output = np.array(bboxes)[indices, :].astype('int'), np.array(confidences)[indices], class_ids[indices]
70-
return output
59+
return output

0 commit comments

Comments
 (0)