if use_cuda: model.cuda() img = cv2.imread(imgfile) # Inference input size is 416*416 does not mean training size is the same # Training size could be 608*608 or even other sizes # Optional inference sizes: # Hight in {320, 416, 512, 608, ... 320 + 96 * n} # Width in {320, 416, 512, 608, ... 320 + 96 * m} sized = cv2.resize(img, (width, height)) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) from tool.utils import load_class_names, plot_boxes_cv2 from tool.torch_utils import do_detect for i in range(2): # This 'for' loop is for speed check # Because the first iteration is usually longer boxes = do_detect(model, sized, 0.4, 0.6, use_cuda) if namesfile == None: if n_classes == 20: namesfile = 'data/voc.names' elif n_classes == 80: namesfile = 'data/coco.names' else: print("please give namefile") class_names = load_class_names(namesfile) plot_boxes_cv2(img, boxes[0], 'predictions.jpg', class_names)
def inference_yolov4(is_local=False): namesfile = None if len(sys.argv) == 6: n_classes = int(sys.argv[1]) weightfile = sys.argv[2] imgfile = sys.argv[3] height = int(sys.argv[4]) width = int(sys.argv[5]) elif len(sys.argv) == 7: n_classes = int(sys.argv[1]) weightfile = sys.argv[2] imgfile = sys.argv[3] height = sys.argv[4] width = int(sys.argv[5]) namesfile = int(sys.argv[6]) else: print('use default value ') #print(' python models.py num_classes weightfile imgfile namefile') n_classes = 2 #WORK_FOLDER = os.path.dirname(os.path.abspath(__file__)) if is_local: WORK_FOLDER = '/fuel/fueling/perception/emergency_detection' else: WORK_FOLDER = '/mnt/bos/modules/perception/emergency_detection' #weightfile = os.path.join(WORK_FOLDER, 'pretrained_model/yolov4.pth') weightfile = os.path.join(WORK_FOLDER, 'checkpoints/Yolov4_epoch600.pth') #weightfile = 'pretrain_model/yolov4.pth' #imgfile = os.path.join(WORK_FOLDER, 'data/dog.jpg') height = 320 width = 320 model = Yolov4(yolov4conv137weight=None, n_classes=n_classes, inference=True) pretrained_dict = torch.load(weightfile, map_location=torch.device('cuda')) model.load_state_dict(pretrained_dict) if is_local: use_cuda = False else: use_cuda = True if use_cuda: model.cuda() if namesfile == None: if n_classes == 20: namesfile = os.path.join(WORK_FOLDER, 'data/voc.names') elif n_classes == 80: namesfile = os.path.join(WORK_FOLDER, 'data/coco.names') elif n_classes == 2: namesfile = os.path.join(WORK_FOLDER, 'data/emergency_vehicle.names') else: print("please give namefile") class_names = load_class_names(namesfile) #image_files = glob.glob(os.path.join(WORK_FOLDER, 'data/emergency_vehicle/images/*.jpg')) csv_file = os.path.join(WORK_FOLDER, 'data/emergency_vehicle/val.csv') with open(csv_file) as f: gt_log = list(csv.reader(f, skipinitialspace=True, delimiter=',', quoting=csv.QUOTE_NONE)) output_file = os.path.join(WORK_FOLDER, 'data/emergency_vehicle/tmp.txt') f = open(output_file, "w") statistics_array = [0,0,0,0,0,0] row_id = 0 for row in gt_log: img_file, ev_label = row img_path = os.path.join(WORK_FOLDER, 'data/emergency_vehicle/images/', img_file) print('processing ', img_path, '...') img = cv2.imread(img_path) # Inference input size is 416*416 does not mean training size is the same # Training size could be 608*608 or even other sizes # Optional inference sizes: # Hight in {320, 416, 512, 608, ... 320 + 96 * n} # Width in {320, 416, 512, 608, ... 320 + 96 * m} sized = cv2.resize(img, (width, height)) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) #for i in range(2): # This 'for' loop is for speed check # Because the first iteration is usually longer boxes = do_detect(model, sized, 0.4, 0.6, use_cuda) #plot_boxes_cv2(img, boxes[0], img_path.replace('images', 'predictions'), class_names) label_string, result_type = get_label_string(img, boxes[0], ev_label, class_names) row = img_file + ' ' + label_string + '\n' if label_string != '': f.write(row) row_id += 1 print(row_id, '/', len(gt_log),': ', row) statistics_array[result_type] += 1 print('************************************************************************') f.close() tot_ev = statistics_array[0]+statistics_array[1]+statistics_array[2] if tot_ev == 0: tot_ev = 1 print('total emergency vehicle ', tot_ev) print('number of correct: ', statistics_array[0], ' wrong: ', statistics_array[1], ' unknown: ', statistics_array[2]) print('ratio of correct: ', statistics_array[0]/tot_ev, ' wrong: ', statistics_array[1]/tot_ev, ' unknown: ', statistics_array[2]/tot_ev) tot_gv = statistics_array[3]+statistics_array[4]+statistics_array[5] if tot_gv == 0: tot_gv = 1 print('total general vehicle ', tot_gv) print('number of correct: ', statistics_array[3], ' wrong: ', statistics_array[4], ' unknown: ', statistics_array[5]) print('ratio of correct: ', statistics_array[3]/tot_gv, ' wrong: ', statistics_array[4]/tot_gv, ' unknown: ', statistics_array[5]/tot_gv)
img = cv2.imread(imgfile) # Inference input size is 416*416 does not mean training size is the same # Training size could be 608*608 or even other sizes # Optional inference sizes: # Hight in {320, 416, 512, 608, ... 320 + 96 * n} # Width in {320, 416, 512, 608, ... 320 + 96 * m} sized = cv2.resize(img, (width, height)) sized = cv2.cvtColor(sized, cv2.COLOR_BGR2RGB) from tool.utils import load_class_names, plot_boxes_cv2 from tool.torch_utils import do_detect for i in range(2): # This 'for' loop is for speed check # Because the first iteration is usually longer boxes = do_detect(model, sized, 0.4, 0.6, use_cuda) """ if namesfile == None: if n_classes == 20: namesfile = 'data/voc.names' elif n_classes == 80: namesfile = 'data/coco.names' else: print("please give namefile") """ class_names = load_class_names('data/kthfs_data/classes.txt') plot_boxes_cv2(img, boxes[0], 'predictions.jpg', class_names)
def demo_tensorflow(tfpb_file="./weight/yolov4.pb", image_path=None, print_sensor_name=False): graph_name = 'yolov4' tf.compat.v1.disable_eager_execution() with tf.compat.v1.Session() as persisted_sess: print("loading graph...") with gfile.FastGFile(tfpb_file, 'rb') as f: graph_def = tf.compat.v1.GraphDef() graph_def.ParseFromString(f.read()) persisted_sess.graph.as_default() tf.import_graph_def(graph_def, name=graph_name) # print all sensor_name if print_sensor_name: tensor_name_list = [ tensor.name for tensor in tf.compat.v1.get_default_graph().as_graph_def().node ] for tensor_name in tensor_name_list: print(tensor_name) inp = persisted_sess.graph.get_tensor_by_name(graph_name + '/' + 'input:0') print(inp.shape) out1 = persisted_sess.graph.get_tensor_by_name(graph_name + '/' + 'output_1:0') out2 = persisted_sess.graph.get_tensor_by_name(graph_name + '/' + 'output_2:0') out3 = persisted_sess.graph.get_tensor_by_name(graph_name + '/' + 'output_3:0') print(out1.shape, out2.shape, out3.shape) # image_src = np.random.rand(1, 3, 608, 608).astype(np.float32) # input image # Input image_src = cv2.imread(image_path) resized = cv2.resize(image_src, (inp.shape[2], inp.shape[3]), interpolation=cv2.INTER_LINEAR) img_in = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB) img_in = np.transpose(img_in, (2, 0, 1)).astype(np.float32) img_in = np.expand_dims(img_in, axis=0) img_in /= 255.0 print("Shape of the network input: ", img_in.shape) feed_dict = {inp: img_in} outputs = persisted_sess.run([out1, out2, out3], feed_dict) print(outputs[0].shape) print(outputs[1].shape) print(outputs[2].shape) boxes = post_processing(img_in, 0.4, outputs) num_classes = 80 if num_classes == 20: namesfile = 'data/voc.names' elif num_classes == 80: namesfile = 'data/coco.names' else: namesfile = 'data/names' class_names = load_class_names(namesfile) result = plot_boxes_cv2(image_src, boxes, savename=None, class_names=class_names) cv2.imshow("tensorflow predicted", result) cv2.waitKey()