import cv2 import numpy as np import os import sys # Disclaimer: majority of the basic framework in this file is modified from the following tutorial: # https://www.learnopencv.com/deep-learning-based-object-detection-using-yolov3-with-opencv-python-c/ # Get the names of the output layers def get_outputs_names(net): # Get the names of all the layers in the network layers_names = net.getLayerNames() # Get the names of the output layers, i.e. the layers with unconnected outputs return [layers_names[i[0] - 1] for i in net.getUnconnectedOutLayers()] # Remove the bounding boxes with low confidence using non-maxima suppression def postprocess(frame, outs, classes, thresh_conf, thresh_nms): frame_height = frame.shape[0] frame_width = frame.shape[1] # Scan through all the bounding boxes output from the network and keep only the # ones with high confidence scores. Assign the box's class label as the class with the highest score. class_ids = [] confidences = [] boxes = [] for out in outs: for detection in out: scores = detection[5:] class_id = np.argmax(scores) confidence = scores[class_id] if confidence > thresh_conf: center_x = int(detection[0] * frame_width) center_y = int(detection[1] * frame_height) width = int(detection[2] * frame_width) height = int(detection[3] * frame_height) left = int(center_x - width / 2) top = int(center_y - height / 2) class_ids.append(class_id) confidences.append(float(confidence)) boxes.append([left, top, width, height]) # Perform non maximum suppression to eliminate redundant overlapping boxes with # lower confidences. indices = cv2.dnn.NMSBoxes(boxes, confidences, thresh_conf, thresh_nms) for i in indices: i = i[0] box = boxes[i] left = box[0] top = box[1] width = box[2] height = box[3] draw_pred(frame, class_ids[i], classes, confidences[i], left, top, left + width, top + height) # Draw the predicted bounding box def draw_pred(frame, class_id, classes, conf, left, top, right, bottom): # Draw a bounding box. cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255)) label = '%.2f' % conf # Get the label for the class name and its confidence if classes: assert (class_id < len(classes)) label = '%s:%s' % (classes[class_id], label) # Display the label at the top of the bounding box label_size, base_line = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) top = max(top, label_size[1]) cv2.putText(frame, label, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255)) def detect_frame(net, classes, img, thresh_conf=0.5, thresh_nms=0.4, in_dim=(416, 416), out_path=None): # Create a 4D blob from a frame. blob = cv2.dnn.blobFromImage(img, 1 / 255, in_dim, [0, 0, 0], 1, crop=False) # Sets the input to the network net.setInput(blob) # Runs the forward pass to get output of the output layers outs = net.forward(get_outputs_names(net)) # Remove the bounding boxes with low confidence postprocess(img, outs, classes, thresh_conf, thresh_nms) # Put efficiency information. The function getPerfProfile returns the # overall time for inference(t) and the timings for each of the layers(in layersTimes) t, _ = net.getPerfProfile() label = 'Inference time: %.2f ms' % (t * 1000.0 / cv2.getTickFrequency()) cv2.putText(img, label, (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255)) if out_path is not None: cv2.imwrite(out_path, img.astype(np.uint8)) def detect_video(net, classes, capture, thresh_conf=0.5, thresh_nms=0.4, in_dim=(416, 416), out_path=None): if out_path is not None: vid_writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 30, (round(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))) while True: ret, frame = capture.read() if not ret: # End of video print("End of video. Press any key to exit") cv2.waitKey(0) break ''' # Create a 4D blob from a frame. blob = cv2.dnn.blobFromImage(frame, 1 / 255, in_dim, [0, 0, 0], 1, crop=False) # Sets the input to the network net.setInput(blob) # Runs the forward pass to get output of the output layers outs = net.forward(get_outputs_names(net)) # Remove the bounding boxes with low confidence postprocess(frame, outs, classes, thresh_conf, thresh_nms) # Put efficiency information. The function getPerfProfile returns the # overall time for inference(t) and the timings for each of the layers(in layersTimes) t, _ = net.getPerfProfile() label = 'Inference time: %.2f ms' % (t * 1000.0 / cv2.getTickFrequency()) cv2.putText(frame, label, (0, 15), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255)) ''' detect_frame(net, classes, frame, thresh_conf=thresh_conf, thresh_nms=thresh_nms, in_dim=in_dim, out_path=None) cv2.imshow('result', frame) if out_path is not None: vid_writer.write(frame.astype(np.uint8)) cv2.waitKey(1) if out_path is not None: vid_writer.release() cv2.destroyAllWindows() pass def main(): # Specify paths for all necessary files test_path = '../data/test1.mp4' weight_path = 'weights/second_general/tiny_yolo_final.weights' cfg_path = 'cfg/tiny_yolo.cfg' class_path = "data/obj.names" out_dir = 'out' if not os.path.isfile(test_path): print('The test file %s doesn\'t exist!' % os.path.abspath(test_path)) if not os.path.isfile(weight_path): print('The weight file %s doesn\'t exist!' % os.path.abspath(test_path)) if not os.path.isfile(cfg_path): print('The config file %s doesn\'t exist!' % os.path.abspath(test_path)) if not os.path.isfile(class_path): print('The class file %s doesn\'t exist!' % os.path.abspath(test_path)) # Setup # Read class names from text file with open(class_path, 'r') as f: classes = [line.strip() for line in f.readlines()] # Load up the neural net using the config and weights net = cv2.dnn.readNetFromDarknet(cfg_path, weight_path) #net.setPreferableBackend(cv2.dnn.DNN_BACKEND_OPENCV) #net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU) # Save the detection result if out_dir is provided if out_dir is None or out_dir == '': out_path = None else: out_path = out_dir + '/' + os.path.split(test_path)[1] # Check if test file is image or video test_ext = test_path[test_path.find('.') + 1:] if test_ext in ['jpg', 'jpeg', 'bmp', 'png', 'tiff']: img = cv2.imread(test_path) detect_frame(net, classes, img, out_path=out_path) else: capture = cv2.VideoCapture(test_path) detect_video(net, classes, capture, out_path=out_path) capture.release() pass if __name__ == '__main__': main()