def save_tflite():
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_tiny(model, FLAGS.weights)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.weights)
model.summary()
converter = tf.lite.TFLiteConverter.from_keras_model(model)
if tf.__version__ >= '2.2.0':
converter.experimental_new_converter = False
if FLAGS.quantize_mode == 'int8':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
elif FLAGS.quantize_mode == 'float16':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_types = [
tf.compat.v1.lite.constants.FLOAT16
]
elif FLAGS.quantize_mode == 'full_int8':
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS_INT8
]
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
converter.allow_custom_ops = True
converter.representative_dataset = representative_data_gen
tflite_model = converter.convert()
open(FLAGS.output, 'wb').write(tflite_model)
logging.info("model saved to: {}".format(FLAGS.output))
def save_tf():
NUM_CLASSES = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASSES)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_tiny(model, FLAGS.weights)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.weights)
model.summary()
model.save(FLAGS.output)
def main(_argv):
print('Arguments', _argv)
print('Flags', flags)
FLAGS.tiny = False
print('Tiny ', FLAGS.tiny)
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
XYSCALE = cfg.YOLO.XYSCALE_TINY
else:
STRIDES = np.array(cfg.YOLO.STRIDES)
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)
XYSCALE = cfg.YOLO.XYSCALE
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_size = FLAGS.size
image_path = FLAGS.image
original_image = cv2.imread(image_path)
print('image:', original_image)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
utils.load_weights_tiny(model, FLAGS.weights, FLAGS.model)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) -
1] == "weights":
utils.load_weights(model, FLAGS.weights)
else:
model.load_weights(FLAGS.weights).expect_partial()
model.summary()
pred_bbox = model.predict(image_data)
else:
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)
interpreter.set_tensor(input_details[0]['index'], image_data)
interpreter.invoke()
pred_bbox = [
interpreter.get_tensor(output_details[i]['index'])
for i in range(len(output_details))
]
if FLAGS.model == 'yolov4':
if FLAGS.tiny:
pred_bbox = utils.postprocess_bbbox(pred_bbox,
ANCHORS,
STRIDES,
XYSCALE,
RESIZE=1.5)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
image = utils.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
#image.show()
print('Image path', image_path)
print('Type Image path', type(image_path))
print('Bboxes type', type(bboxes))
classes = utils.read_class_names(cfg.YOLO.CLASSES)
list_bboxes = []
for i, bbox in enumerate(bboxes):
coor = np.array(bbox[:4], dtype=np.int32)
score = bbox[4]
class_ind = int(bbox[5])
#print('type bbox',type(bbox))
#print('bbox',bbox[:4])
#print('coor',list(coor))
bbox_info = {
'coor': list(coor),
'probability': score,
'class': classes[class_ind]
}
list_bboxes.append(bbox_info)
try:
output_name = os.path.join('results/out_' +
os.path.basename(image_path))
image.save(output_name)
#cv2.imwrite(output_name,img)
print('Img saved to', output_name)
output = pd.DataFrame(list_bboxes)
print('image_path', image_path)
output_name = '.'.join(output_name.split('.')[:2]) + '.xlsx'
#output_name = 'results/out_'+image_path.split('\\')[-1].split('.')[0]+'.xlsx'
print('output_name', output_name)
output.to_excel(output_name)
except Exception as e:
print(e)
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
XYSCALE = cfg.YOLO.XYSCALE_TINY
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY_V3, FLAGS.tiny)
else:
STRIDES = np.array(cfg.YOLO.STRIDES)
XYSCALE = cfg.YOLO.XYSCALE
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)
CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)
NUM_CLASSES = len(CLASSES)
input_size = FLAGS.size
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
times = []
if FLAGS.output:
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps, (width, height))
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASSES)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_tiny(model, FLAGS.weights, FLAGS.model)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) -
1] == "weights":
utils.load_weights(model, FLAGS.weights)
else:
model.load_weights(FLAGS.weights).expect_partial()
model.summary()
elif FLAGS.framework == 'tflite':
interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
elif FLAGS.framework == 'trt':
saved_model_loaded = tf.saved_model.load(FLAGS.weights,
tags=[tag_constants.SERVING])
infer = saved_model_loaded.signatures['serving_default']
max_cosine_distance = 0.7 # 0.5 / 0.7
nn_budget = None
model_filename = './weights/tracker/mars-small128.pb'
encoder = gdet.create_box_encoder(model_filename, batch_size=1)
metric = nn_matching.NearestNeighborDistanceMetric("cosine",
max_cosine_distance,
nn_budget)
tracker = Tracker(metric)
key_list = list(CLASSES.keys())
val_list = list(CLASSES.values())
Track_only = []
logging.info("Models loaded!")
while True:
return_value, frame = vid.read()
if not return_value:
logging.warning("Empty Frame")
break
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame_size = frame.shape[:2]
image_data = utils.image_preprocess(np.copy(frame),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
t1 = time.time()
if FLAGS.framework == 'tf':
pred_bbox = model.predict(image_data)
elif FLAGS.framework == 'tflite':
interpreter.set_tensor(input_details[0]['index'], image_data)
interpreter.invoke()
pred_bbox = [
interpreter.get_tensor(output_details[i]['index'])
for i in range(len(output_details))
]
elif FLAGS.framework == 'trt':
batched_input = tf.constant(image_data)
pred_bbox = []
result = infer(batched_input)
for _, value in result.items():
value = value.numpy()
pred_bbox.append(value)
t2 = time.time()
times.append(t2 - t1)
times = times[-20:]
ms = sum(times) / len(times) * 1000
fps = 1000 / ms
if FLAGS.model == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size,
0.5) # 0.25
bboxes = utils.nms(bboxes, 0.5, method='nms') # 0.213
boxes, scores, names = [], [], []
for bbox in bboxes:
if len(Track_only) != 0 and CLASSES[int(
bbox[5])] in Track_only or len(Track_only) == 0:
boxes.append([
bbox[0].astype(int), bbox[1].astype(int),
bbox[2].astype(int) - bbox[0].astype(int),
bbox[3].astype(int) - bbox[1].astype(int)
])
scores.append(bbox[4])
names.append(CLASSES[int(bbox[5])])
boxes = np.array(boxes)
names = np.array(names)
scores = np.array(scores)
features = np.array(encoder(frame, boxes))
detections = [
Detection(bbox, score, class_name, feature)
for bbox, score, class_name, feature in zip(
boxes, scores, names, features)
]
tracker.predict()
tracker.update(detections)
tracked_bboxes = []
for track in tracker.tracks:
if not track.is_confirmed(
) or track.time_since_update > 1: # 1 / 5
continue
bbox = track.to_tlbr()
class_name = track.get_class()
tracking_id = track.track_id
index = key_list[val_list.index(class_name)]
tracked_bboxes.append(bbox.tolist() + [tracking_id, index])
image = utils.draw_bbox(frame,
tracked_bboxes,
classes=CLASSES,
tracking=True)
image = cv2.putText(
image,
"Time: {:.2f}ms".format(sum(times) / len(times) * 1000),
(0, 36), # 24
cv2.FONT_HERSHEY_SIMPLEX,
1.5,
(0, 0, 255),
2)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
cv2.namedWindow("Detections", cv2.WINDOW_AUTOSIZE)
cv2.imshow("Detections", image)
if FLAGS.output:
out.write(image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
vid.release()
if FLAGS.output:
out.release()
cv2.destroyAllWindows()
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
XYSCALE = cfg.YOLO.XYSCALE_TINY
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY_V3, FLAGS.tiny)
else:
STRIDES = np.array(cfg.YOLO.STRIDES)
XYSCALE = cfg.YOLO.XYSCALE
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)
CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)
NUM_CLASSES = len(CLASSES)
input_size = FLAGS.size
try:
vid = cv2.VideoCapture(int(FLAGS.video))
except:
vid = cv2.VideoCapture(FLAGS.video)
times = []
if FLAGS.output:
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = int(vid.get(cv2.CAP_PROP_FPS))
codec = cv2.VideoWriter_fourcc(*FLAGS.output_format)
out = cv2.VideoWriter(FLAGS.output, codec, fps,
(width, height)) # TODO: switch to get vertical
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASSES)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_tiny(model, FLAGS.weights, FLAGS.model)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) -
1] == "weights":
utils.load_weights(model, FLAGS.weights)
else:
model.load_weights(FLAGS.weights).expect_partial()
model.summary()
elif FLAGS.framework == 'tflite':
interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
elif FLAGS.framework == 'trt':
saved_model_loaded = tf.saved_model.load(FLAGS.weights,
tags=[tag_constants.SERVING])
infer = saved_model_loaded.signatures['serving_default']
logging.info("Model loaded!")
while True:
return_value, frame = vid.read()
# frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE) # TODO: here
if not return_value:
logging.warning("Empty Frame")
break
frame_size = frame.shape[:2]
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image_data = utils.image_preprocess(np.copy(frame),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
prev_time = time.time()
if FLAGS.framework == 'tf':
pred_bbox = model.predict(image_data)
elif FLAGS.framework == 'tflite':
interpreter.set_tensor(input_details[0]['index'], image_data)
interpreter.invoke()
pred_bbox = [
interpreter.get_tensor(output_details[i]['index'])
for i in range(len(output_details))
]
elif FLAGS.framework == 'trt':
batched_input = tf.constant(image_data)
pred_bbox = []
result = infer(batched_input)
for _, value in result.items():
value = value.numpy()
pred_bbox.append(value)
curr_time = time.time()
times.append(curr_time - prev_time)
times = times[-20:]
if FLAGS.model == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size,
0.5) # 0.25
bboxes = utils.nms(bboxes, 0.213, method='nms') # 0.213
image = utils.draw_bbox(frame, bboxes, classes=CLASSES)
image = cv2.putText(
image,
"Time: {:.2f}ms".format(sum(times) / len(times) * 1000),
(0, 24), # 24
cv2.FONT_HERSHEY_SIMPLEX,
0.7,
(0, 0, 255),
2) # 0.7
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.namedWindow("Detections", cv2.WINDOW_AUTOSIZE)
cv2.imshow("Detections", image)
if FLAGS.output:
out.write(image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
vid.release()
if FLAGS.output:
out.release()
cv2.destroyAllWindows()
def main(_argv):
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
XYSCALE = cfg.YOLO.XYSCALE_TINY
else:
STRIDES = np.array(cfg.YOLO.STRIDES)
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)
XYSCALE = cfg.YOLO.XYSCALE
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_size = FLAGS.size
image_path = FLAGS.image
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image), [input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
utils.load_weights_tiny(model, FLAGS.weights, FLAGS.model)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
utils.load_weights(model, FLAGS.weights)
else:
model.load_weights(FLAGS.weights).expect_partial()
model.summary()
pred_bbox = model.predict(image_data)
else:
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)
interpreter.set_tensor(input_details[0]['index'], image_data)
interpreter.invoke()
pred_bbox = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))]
if FLAGS.model == 'yolov4':
if FLAGS.tiny:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES, XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES, XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size, input_size, 0.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
image = utils.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
image.show()
image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
cv2.imwrite(FLAGS.output, image)
def __init__(self, framework='tf', size=608, tiny=False, model='yolov4'):
self.framework = framework
self.weights = 'weights' #None#weights
self.size = size
self.tiny = tiny
self.model = model
self.instanciated_model = None
# Instanciate model
print('Tiny ', self.tiny)
#image_path = self.image
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_size = self.size
if self.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if self.tiny:
if self.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
utils.load_weights_tiny(model, self.weights, self.model)
else:
if self.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, self.weights)
elif self.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
# Check if files have already been downloaded
yolov4_weights_path = os.path.join(self.weights,
'yolov4.weights')
#yolov4_weights_path = os.path.join(Path(os.path.realpath(__file__)).parent,'data/yolov4.weights')
if not os.path.exists(yolov4_weights_path):
print('Downloading weights file')
self.weights = self.download('yolov4.weights',
local_path=self.weights)
print('Weight file was downloaded to', self.weights)
else:
print('Weights file already downloaded')
self.weights = yolov4_weights_path
if self.weights.split(".")[len(self.weights.split(".")) -
1] == "weights":
utils.load_weights(model, self.weights)
else:
model.load_weights(self.weights).expect_partial()
self.instanciated_model = model
else:
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path=self.weights)
interpreter.allocate_tensors()
self.instanciated_model = interpreter
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
XYSCALE = cfg.YOLO.XYSCALE_TINY
else:
STRIDES = np.array(cfg.YOLO.STRIDES)
XYSCALE = cfg.YOLO.XYSCALE
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
input_size = FLAGS.size
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASS)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_tiny(model, FLAGS.weights)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.weights)
elif FLAGS.framework == 'trt':
saved_model_loaded = tf.saved_model.load(FLAGS.weights,
tags=[tag_constants.SERVING])
signature_keys = list(saved_model_loaded.signatures.keys())
# print(signature_keys)
infer = saved_model_loaded.signatures['serving_default']
logging.info('weights loaded')
@tf.function
def run_model(x):
return model(x)
# Test the TensorFlow Lite model on random input data.
sum = 0
original_image = cv2.imread(FLAGS.image)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image),
[FLAGS.size, FLAGS.size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
img_raw = tf.image.decode_image(open(FLAGS.image, 'rb').read(), channels=3)
img_raw = tf.expand_dims(img_raw, 0)
img_raw = tf.image.resize(img_raw, (FLAGS.size, FLAGS.size))
batched_input = tf.constant(image_data)
for i in range(1000):
prev_time = time.time()
# pred_bbox = model.predict(image_data)
if FLAGS.framework == 'tf':
pred_bbox = []
result = run_model(image_data)
for value in result:
value = value.numpy()
pred_bbox.append(value)
if FLAGS.model == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS,
STRIDES, XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS,
STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
elif FLAGS.framework == 'trt':
pred_bbox = []
result = infer(batched_input)
for key, value in result.items():
value = value.numpy()
pred_bbox.append(value)
if FLAGS.model == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS,
STRIDES, XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS,
STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
# pred_bbox = pred_bbox.numpy()
curr_time = time.time()
exec_time = curr_time - prev_time
if i == 0:
continue
sum += (1 / exec_time)
info = str(i) + " time:" + str(round(
exec_time, 3)) + " average FPS:" + str(round(
sum / i, 2)) + ", FPS: " + str(round((1 / exec_time), 1))
print(info)
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
for physical_device in physical_devices:
tf.config.experimental.set_memory_growth(physical_device, True)
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
XYSCALE = cfg.YOLO.XYSCALE_TINY
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY_V3, FLAGS.tiny)
else:
STRIDES = np.array(cfg.YOLO.STRIDES)
XYSCALE = cfg.YOLO.XYSCALE
if FLAGS.model == 'yolov4':
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS, FLAGS.tiny)
else:
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_V3, FLAGS.tiny)
CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)
NUM_CLASSES = len(CLASSES)
input_size = FLAGS.size
image_path = FLAGS.image
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
image_data = utils.image_preprocess(np.copy(original_image),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3_tiny(input_layer, NUM_CLASSES)
else:
feature_maps = YOLOv4_tiny(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_tiny(model, FLAGS.weights, FLAGS.model)
else:
if FLAGS.model == 'yolov3':
feature_maps = YOLOv3(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights_v3(model, FLAGS.weights)
elif FLAGS.model == 'yolov4':
feature_maps = YOLOv4(input_layer, NUM_CLASSES)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) -
1] == "weights":
utils.load_weights(model, FLAGS.weights)
else:
model.load_weights(FLAGS.weights).expect_partial()
model.summary()
pred_bbox = model.predict(image_data)
elif FLAGS.framework == 'tflite':
interpreter = tf.lite.Interpreter(model_path=FLAGS.weights)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
interpreter.set_tensor(input_details[0]['index'], image_data)
interpreter.invoke()
pred_bbox = [
interpreter.get_tensor(output_details[i]['index'])
for i in range(len(output_details))
]
elif FLAGS.framework == 'trt':
saved_model_loaded = tf.saved_model.load(FLAGS.weights,
tags=[tag_constants.SERVING])
infer = saved_model_loaded.signatures['serving_default']
batched_input = tf.constant(image_data)
pred_bbox = []
result = infer(batched_input)
for _, value in result.items():
value = value.numpy()
pred_bbox.append(value)
if FLAGS.model == 'yolov4':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
else:
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, 0.5) # 0.25
bboxes = utils.nms(bboxes, 0.5, method='nms') # 0.213
image = utils.draw_bbox(original_image, bboxes, classes=CLASSES)
image = cv2.cvtColor(np.array(image), cv2.COLOR_BGR2RGB)
cv2.imwrite(FLAGS.output, image)