How to count objectsΒΆ
Here is the code related to counting objects for YOLOv4
In the main while True: loop:
if FLAGS.count:
# count objects found
counted_classes = utils.count_objects(
pred_bbox, by_class=True, allowed_classes=allowed_classes)
# draw bounding boxes and output count
image = utils.draw_bbox(
frame, pred_bbox, FLAGS.info, counted_classes, allowed_classes=allowed_classes)
else:
# just draw the bounding boxes
image = utils.draw_bbox(
frame, pred_bbox, FLAGS.info, allowed_classes=allowed_classes)
Here is utils.count_objects():
# function to count objects, can return total classes or count per class
def count_objects(data, by_class=False, allowed_classes=list(read_class_names(cfg.YOLO.CLASSES).values())):
boxes, scores, classes, num_objects = data
#create dictionary to hold count of objects
counts = dict()
# if by_class = True then count objects per class
if by_class:
class_names = read_class_names(cfg.YOLO.CLASSES)
# loop through total number of objects found
for i in range(num_objects):
# grab class index and convert into corresponding class name
class_index = int(classes[i])
class_name = class_names[class_index]
if class_name in allowed_classes:
counts[class_name] = counts.get(class_name, 0) + 1
else:
continue
# else count total objects found
else:
counts['total object'] = num_objects
return counts
In the draw_bbox function:
if counted_classes != None:
for key, value in counted_classes.items():
if key == "bench" or key == "bus" or key == "wine glass" or key == "sandwich" or key == "toothbrush":
cv2.putText(image, "{}es detected: {}".format(key, value), (indent, offset),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, red, font_thickness)
elif key == "sheep" or key == "skis" or key == "scissors":
cv2.putText(image, "{} detected: {}".format(key, value), (indent, offset),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, red, font_thickness)
elif key == "knife":
cv2.putText(image, "knives detected: {}".format(value), (indent, offset),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, red, font_thickness)
elif key == "mouse":
cv2.putText(image, "mice detected: {}".format(value), (indent, offset),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, red, font_thickness)
elif key == "person":
cv2.putText(image, "people detected: {}".format(value), (indent, offset),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, blue, font_thickness)
else:
cv2.putText(image, "{}s detected: {}".format(key, value), (indent, offset),
cv2.FONT_HERSHEY_COMPLEX_SMALL, font_size, red, font_thickness)
offset += 40
return image