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()
