def save_tf():
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
feature_maps = YOLOv4(input_layer, NUM_CLASS)
bbox_tensors = []
prob_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE)
elif i == 1:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
boxes, pred_conf = filter_boxes(pred_bbox,
pred_prob,
score_threshold=FLAGS.score_thres,
input_shape=tf.constant(
[FLAGS.input_size, FLAGS.input_size]))
pred = tf.concat([boxes, pred_conf], axis=-1)
model = tf.keras.Model(input_layer, pred)
utils.load_weights(model, FLAGS.weights, FLAGS.model)
model.summary()
model.save(FLAGS.output)
def main(_argv):
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:
feature_maps = YOLOv3_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)
save_model(model)
save_tflite(model)
demo()
def my_decode(feature_maps):
global NUM_CLASS, STRIDES, ANCHORS, XYSCALE
# pred_xywh, pred_conf, pred_prob
#res = []
bbox_tensors = []
prob_tensors = []
for i, fm in enumerate(feature_maps):
with tf.name_scope("featuremap-" + str(i)) as scope:
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
#output_tensors = decode_tf(fm, FLAGS.input_size // 8, NUM_CLASS, STRIDES, ANCHORS)
#output_tensors = decode_train(fm, FLAGS.input_size // 8, NUM_CLASS, STRIDES, ANCHORS)
elif i == 1:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
#output_tensors = decode_tf(fm, FLAGS.input_size // 16, NUM_CLASS, STRIDES, ANCHORS)
#output_tensors = decode_train(fm, FLAGS.input_size // 16, NUM_CLASS, STRIDES, ANCHORS)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
#output_tensors = decode_tf(fm, FLAGS.input_size // 32, 8UM_CLASS, STRIDES, ANCHORS)
#output_tensors = decode_train(fm, FLAGS.input_size // 32, NUM_CLASS, STRIDES, ANCHORS)
#res.append(output_tensors)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
#pred = res
pred = (pred_bbox, pred_prob)
return pred
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:
feature_maps = YOLOv3_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()
# Save Model, including Frozen graph
model.save('./model')
def save_tf():
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:
feature_maps = YOLOv3_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)
else:
print("model option can be only 'yolov3' or 'yolov4'.")
return
model.summary()
model.save(FLAGS.output)
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 main(_argv):
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_size = FLAGS.size
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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)
logging.info('weights loaded')
# Test the TensorFlow Lite model on random input data.
sum = 0
for i in range(1000):
img_raw = tf.image.decode_image(open(FLAGS.image, 'rb').read(),
channels=3)
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_preporcess(np.copy(original_image),
[FLAGS.size, FLAGS.size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
prev_time = time.time()
pred_bbox = model.predict(image_data)
# pred_bbox = pred_bbox.numpy()
curr_time = time.time()
exec_time = curr_time - prev_time
if i == 0: continue
sum += (1000 / (1000 * exec_time))
info = "average FPS:" + str(round(sum / i, 2)) + ", FPS: " + str(
round((1000 / (1000 * exec_time)), 1))
print(info)
def save_tf():
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
bbox_tensors = []
prob_tensors = []
if FLAGS.tiny:
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
else:
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
elif i == 1:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
if FLAGS.framework == 'tflite':
pred = (pred_bbox, pred_prob)
else:
boxes, pred_conf = filter_boxes(
pred_bbox,
pred_prob,
score_threshold=FLAGS.score_thres,
input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
pred = tf.concat([boxes, pred_conf], axis=-1)
model = tf.keras.Model(input_layer, pred)
model.load_weights(FLAGS.weights)
#utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny) #weight파일일 경
#model.summary()
#model.save('/checkpoints/yolov4-416')
#model.save(FLAGS.output, save_format = 'tf')
tf.saved_model.save(
model, FLAGS.output
) #현재 이 저장 형태의 경우: assets/, variables/, saved_model.pb 형태로 저장됨.
def save_tf():
if FLAGS.license:
cfg.YOLO.CLASSES = "./data/classes/custom.names"
else:
cfg.YOLO.CLASSES = "./data/classes/char.names"
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
#print(read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
bbox_tensors = []
prob_tensors = []
if FLAGS.tiny:
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
else:
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
elif i == 1:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
if FLAGS.framework == 'tflite':
pred = (pred_bbox, pred_prob)
else:
boxes, pred_conf = filter_boxes(
pred_bbox,
pred_prob,
score_threshold=FLAGS.score_thres,
input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
pred = tf.concat([boxes, pred_conf], axis=-1)
model = tf.keras.Model(input_layer, pred)
utils.load_weights(model, FLAGS.weights, FLAGS.model, FLAGS.tiny)
model.summary()
model.save(FLAGS.output)
def save_tf():
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model, FLAGS.tiny)
bbox_tensors = []
prob_tensors = []
if FLAGS.tiny:
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
else:
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
elif i == 1:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
if FLAGS.framework == 'tflite':
pred = (pred_bbox, pred_prob)
else:
boxes, pred_conf = filter_boxes(
pred_bbox,
pred_prob,
score_threshold=FLAGS.score_thres,
input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
pred = tf.concat([boxes, pred_conf], axis=-1)
model = tf.keras.Model(input_layer, pred)
model.load_weights(FLAGS.input_model_path)
model.summary()
# model.save(FLAGS.output_model_path)
return model
def save_tf(parameters):
"""Transform a darknet model of YOLO to a TensorFlow model
Args:
parameters (dictionary): input parameters
- weights: path to the darknet weights
- input_size: input size of the model
- model: model to transform
- weights_tf: path to save the tf weights
Returns:
[void]:
"""
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(
tiny=False, model=parameters['model'])
input_layer = tf.keras.layers.Input(
[parameters['input_size'], parameters['input_size'], 3])
feature_maps = YOLO(input_layer, NUM_CLASS, parameters['model'], False)
bbox_tensors = []
prob_tensors = []
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, parameters['input_size'] // 8,
NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE,
'tf')
elif i == 1:
output_tensors = decode(fm, parameters['input_size'] // 16,
NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE,
'tf')
else:
output_tensors = decode(fm, parameters['input_size'] // 32,
NUM_CLASS, STRIDES, ANCHORS, i, XYSCALE,
'tf')
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
boxes, pred_conf = filter_boxes(pred_bbox,
pred_prob,
score_threshold=parameters['score_thres'],
input_shape=tf.constant([
parameters['input_size'],
parameters['input_size']
]))
pred = tf.concat([boxes, pred_conf], axis=-1)
model = tf.keras.Model(input_layer, pred)
utils.load_weights(model, parameters['weights'], parameters['model'],
False)
model.summary()
model.save(parameters['weights_tf'])
def start_model(device):
tf.executing_eagerly()
#TODO will have to change when working with several gpus
# strategy = tf.distribute.MirroredStrategy()
with None: #tf.device(device):
#generate model
input_layer = tf.keras.Input([INPUT_SIZE, INPUT_SIZE, 3])
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)
print('Model built')
#force to run eagerly
model.run_eagerly = True
#load existing weights into model
utils.load_weights(model, WEIGHTS)
return model
def main(argv):
weights = FLAGS.weights
input_size = FLAGS.size
NUM_CLASS = 2
print(f'[DEBUG][save_model] Path to weights : weights/{FLAGS.weights}')
print(f'[DEBUG][save_model] Size : {FLAGS.size}')
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
print(f'[INFO][save_model] Created input_layer of size {input_size}')
print(f'[DEBUG][save_model] input_layer : {input_layer}')
feature_maps = YOLOv4(input_layer, NUM_CLASS)
print(f'[DEBUG][save_model] feature_maps : {feature_maps}')
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensors.append(decode(fm, NUM_CLASS, i))
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, 'weights/' + FLAGS.weights)
print(f'[INFO][save_model] Saving model... ')
model.save(f'models/{weights.split(".")[0]}-size-{input_size}.h5')
print(f'[INFO][save_model] Model saved to models/{weights.split(".")[0]}-size-{input_size}.h5')
def main(_argv):
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
input_size = FLAGS.size
video_path = FLAGS.video
print("Video from: ", video_path )
vid = cv2.VideoCapture(video_path)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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)
def save_tf():
STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS)
print("load_config XYSCALE:{}".format(XYSCALE))
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
feature_maps = YOLO(input_layer, NUM_CLASS, FLAGS.model,
FLAGS.num_detection_layer)
bbox_tensors = []
prob_tensors = []
if FLAGS.num_detection_layer == 1: # yolo-custom
output_tensors = decode(feature_maps[0], FLAGS.input_size // 32,
NUM_CLASS, STRIDES, ANCHORS, 0, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
elif FLAGS.num_detection_layer == 2: # yolo-tiny
for i, fm in enumerate(feature_maps):
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
elif FLAGS.num_detection_layer == 3: # yolo
for i, fm in enumerate(feature_maps):
print("i:{}".format(i))
if i == 0:
output_tensors = decode(fm, FLAGS.input_size // 8, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
elif i == 1:
output_tensors = decode(fm, FLAGS.input_size // 16, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
else:
output_tensors = decode(fm, FLAGS.input_size // 32, NUM_CLASS,
STRIDES, ANCHORS, i, XYSCALE,
FLAGS.framework)
bbox_tensors.append(output_tensors[0])
prob_tensors.append(output_tensors[1])
pred_bbox = tf.concat(bbox_tensors, axis=1)
pred_prob = tf.concat(prob_tensors, axis=1)
if FLAGS.framework == 'tflite':
pred = (pred_bbox, pred_prob)
else:
boxes, pred_conf = filter_boxes(
pred_bbox,
pred_prob,
score_threshold=FLAGS.score_thres,
input_shape=tf.constant([FLAGS.input_size, FLAGS.input_size]))
pred = tf.concat([boxes, pred_conf], axis=-1)
model = tf.keras.Model(input_layer, pred)
utils.load_weights(model, FLAGS.weights, FLAGS.model,
FLAGS.num_detection_layer)
model.summary()
model.save(FLAGS.output)
def load_model_yolov4(input_size=608,
weights='./data/yolov4_original_last.weights'):
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
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, weights)
model.summary()
return model
def create_ppe_detector(input_size):
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
feature_maps = YOLOv4(input_layer, NUM_CLASSES_PPE)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASSES_PPE, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.ppe_weights)
model.summary()
return model
elif FLAGS.framework == 'trt':
saved_model_loaded = tf.saved_model.load(FLAGS.ppe_weights,
tags=[tag_constants.SERVING])
infer = saved_model_loaded.signatures['serving_default']
return infer
def on_init():
global model
global anchors
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
# Restrict TensorFlow to only allocate 1GB of memory on the first GPU
try:
tf.config.experimental.set_virtual_device_configuration(
gpus[0], [
tf.config.experimental.VirtualDeviceConfiguration(
memory_limit=1024)
])
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
sys.stdout.write(
f"{len(gpus)} Physical GPUs, {len(logical_gpus)} Logical GPUs\n"
)
sys.stdout.flush()
except RuntimeError as e:
# Virtual devices must be set before GPUs have been initialized
sys.stdout.write(e)
sys.stdout.flush()
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
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 weights.split(".")[len(weights.split(".")) - 1] == "weights":
utils.load_weights(model, weights)
else:
model.load_weights(weights).expect_partial()
if model_name == "yolov3":
anchors = utils.get_anchors(ANCHOR_V3_DEFAULT)
else:
anchors = utils.get_anchors(ANCHOR_V4_DEFAULT)
return True
def main(_argv):
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
input_layer = tf.keras.layers.Input([FLAGS.size, FLAGS.size, 3])
feature_maps = YOLOv4(input_layer)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.weights)
logging.info('weights loaded')
# Test the TensorFlow Lite model on random input data.
for i in range(1000):
img_raw = tf.image.decode_image(open(FLAGS.image, 'rb').read(),
channels=3)
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_preporcess(np.copy(original_image),
[FLAGS.size, FLAGS.size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
prev_time = time.time()
pred_bbox = model.predict(image_data)
# pred_bbox = pred_bbox.numpy()
curr_time = time.time()
exec_time = curr_time - prev_time
info = "time:" + str(round(1000 * exec_time, 2)) + " ms, FPS: " + str(
round((1000 / (1000 * exec_time)), 1))
print(info)
def main(_argv):
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
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_preporcess(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:
feature_maps = YOLOv3_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()
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':
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()
def main():
tf.executing_eagerly()
strategy = tf.distribute.OneDeviceStrategy(device="/gpu:0")
with strategy.scope():
# if True:
STRIDES = np.array(cfg.YOLO.STRIDES)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
WEIGHTS = './data/yolov4.weights' #must end in .weights
video_path = './data/road.mp4'
video_path = './data/AOTsample3.mp4'
#video_path = './data/vtest.avi'
#video_path = './data/20190422_153844_DA4A.mkv'
print("Video from: ", video_path )
#vid = cv2.VideoCapture(video_path)
print('thread started')
INPUT_SIZE = 419 #608 #230
#open file to output to
output_f = video_path[:-3] + 'txt'
f = open(output_f, 'w')
print('file started')
#generate model
input_layer = tf.keras.Input([INPUT_SIZE, INPUT_SIZE, 3])
print('tensors started 1')
feature_maps = YOLOv4(input_layer, NUM_CLASS)
print('tensors started 2')
bbox_tensors = []
print('tensors started 3')
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, NUM_CLASS, i)
bbox_tensors.append(bbox_tensor)
print('tensors started 4')
model = tf.keras.Model(input_layer, bbox_tensors)
print('model built')
#force to run eagerly
model.run_eagerly = True
if model.run_eagerly:
print ('yeeyee')
else:
print ('hawhaw')
utils.load_weights(model, WEIGHTS)
with tf.device('/GPU:0'):
buf = Queue(maxsize=8)
# buf = VidThread(video_path)
# buf.start()
vid = cv2.VideoCapture(video_path)
coord = tf.train.Coordinator()
t = Thread(target=MyLoop, args=(video_path, buf,vid, coord))
t.daemon = True
#coord.register_thread(t)
t.start()
time.sleep(1.0)
try:
while not buf.empty():
frame = buf.get()
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame)
dt = str(datetime.datetime.now())
frame_size = frame.shape[:2]
#resize image and add another dimension
cur_frame = np.copy(frame)
image_data = utils.image_preprocess(cur_frame, [INPUT_SIZE, INPUT_SIZE])
image_data = image_data[np.newaxis, ...].astype(np.float32)
prev_time = time.time()
with tf.device('/GPU:0'):
image_data = tf.convert_to_tensor(image_data)
print(image_data.device)
curr_time = time.time()
exec_time = curr_time - prev_time
info = "time1: %.2f ms" %(1000*exec_time)
print(info)
prev_time = time.time()
#make bboxes
pred_bbox = model.predict(image_data)
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES, XYSCALE)
bboxes = utils.postprocess_boxes(pred_bbox, frame_size, INPUT_SIZE, 0.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
#output bbox info to file and show image
#calculate and display time it took to process frame
utils.video_write_info(frame, f, bboxes, dt)
image = utils.draw_some_bbox(frame, bboxes)
curr_time = time.time()
exec_time = curr_time - prev_time
info = "time2: %.2f ms" %(1000*exec_time)
print(info)
result = np.asarray(image)
cv2.namedWindow("result", cv2.WINDOW_NORMAL)
result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR) #swapped image with result, not sure what the effect was
cv2.imshow("result", result)
if cv2.waitKey(1) & 0xFF == ord('q'): break
#end video, close viewer, stop writing to file
vid.release()
cv2.destroyAllWindows()
f.close()
#if interrupted, end video, close viewer, stop writing to file
except:
print("Unexpected error:", sys.exc_info()[0])
vid.release()
cv2.destroyAllWindows()
f.close()
def save_tflite():
input_layer = tf.keras.layers.Input(
[FLAGS.input_size, FLAGS.input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv4(input_layer)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
utils.load_weights_tiny(model, FLAGS.weights)
else:
feature_maps = YOLOv4(input_layer)
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensor = decode(fm, i)
bbox_tensors.append(bbox_tensor)
model = tf.keras.Model(input_layer, bbox_tensors)
model.summary()
utils.load_weights(model, FLAGS.weights)
converter = tf.lite.TFLiteConverter.from_keras_model(model)
if FLAGS.quantize_mode == 'int8':
converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.default_ranges_stats = (0, 6)
# converter.inference_type = tf.compat.v1.lite.constants.QUANTIZED_UINT8
# converter.output_format = tf.compat.v1.lite.constants.TFLITE
# converter.allow_custom_ops = True
# converter.quantized_input_stats = {"input0": (0., 1.)}
# converter.optimizations = [tf.lite.Optimize.DEFAULT]
# converter.inference_input_type = tf.compat.v1.lite.constants.QUANTIZED_UINT8
# converter.inference_output_type = tf.compat.v1.lite.constants.QUANTIZED_UINT8
# converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
# converter.inference_input_type = tf.uint8
# converter.inference_output_type = tf.uint8
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)
# tflite_model = converter.convert()
# tf.GFile(FLAGS.output, "wb").write(tflite_model)
logging.info("model saved to: {}".format(FLAGS.output))
def main(_argv):
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
config = ConfigProto()
config.gpu_options.allow_growth = True
session = InteractiveSession(config=config)
input_size = FLAGS.size
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
tf.config.experimental.set_memory_growth(physical_devices[0], True)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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):
NUM_CLASS = 2
ANCHORS = [
12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243,
459, 401
]
ANCHORS = np.array(ANCHORS, dtype=np.float32)
ANCHORS = ANCHORS.reshape(3, 3, 2)
STRIDES = [8, 16, 32]
XYSCALE = [1.2, 1.1, 1.05]
input_size = FLAGS.size
image_path = FLAGS.image_path
score_thresh = FLAGS.score_thresh
iou_thresh = FLAGS.iou_thresh
save_path = FLAGS.save_path
print(f'[DEBUG][image] input_size : {input_size}')
print(f'[DEBUG][image] image_path : {image_path}')
print(f'[DEBUG][image] score_thresh : {score_thresh}')
print(f'[DEBUG][image] iou_thresh : {iou_thresh}')
original_image = cv2.imread(image_path)
original_image = cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB)
original_image_size = original_image.shape[:2]
print(f'[DEBUG][image] original_image_size : {original_image_size}')
image_data = utils.image_preprocess(np.copy(original_image),
[input_size, input_size])
image_data = image_data[np.newaxis, ...].astype(np.float32)
print('[INFO] Bulding Yolov4 architecture')
tic = time.perf_counter()
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
print(f'[INFO][image] Created input_layer of size {input_size}')
print(f'[DEBUG][image] input_layer : {input_layer}')
feature_maps = YOLOv4(input_layer, NUM_CLASS)
print(f'[DEBUG][image] feature_maps : {feature_maps}')
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensors.append(decode(fm, NUM_CLASS, i))
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.weights)
toc = time.perf_counter()
print(f'[INFO] Architecture built.')
print(f'[DEBUG][image] Execution took {(1000 * (toc - tic)):0.4f} ms')
pred_bbox = model.predict(image_data)
print(f'[INFO][image] Finished initial predication on image')
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES, XYSCALE)
bboxes = utils.postprocess_boxes(pred_bbox, original_image_size,
input_size, score_thresh)
bboxes = utils.nms(bboxes, iou_thresh, method='nms')
image = utils.draw_bbox(original_image, bboxes)
image = Image.fromarray(image)
image.show()
if (save_path):
image.save(save_path)
print(f'[INFO][image] Detected image saved to {save_path}')
def main(argv):
NUM_CLASS = 2
ANCHORS = [
12, 16, 19, 36, 40, 28, 36, 75, 76, 55, 72, 146, 142, 110, 192, 243,
459, 401
]
ANCHORS = np.array(ANCHORS, dtype=np.float32)
ANCHORS = ANCHORS.reshape(3, 3, 2)
STRIDES = [8, 16, 32]
XYSCALE = [1.2, 1.1, 1.05]
input_size = FLAGS.size
score_thresh = FLAGS.score_thresh
iou_thresh = FLAGS.iou_thresh
save_path = FLAGS.save_path
print(f'[DEBUG][webcam] input_size : {input_size}')
print(f'[DEBUG][webcam] score_thresh : {score_thresh}')
print(f'[DEBUG][webcam] iou_thresh : {iou_thresh}')
print('[INFO] Bulding Yolov4 architecture')
tic = time.perf_counter()
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
print(f'[INFO][webcam] Created input_layer of size {input_size}')
print(f'[DEBUG][webcam] input_layer : {input_layer}')
feature_maps = YOLOv4(input_layer, NUM_CLASS)
print(f'[DEBUG][webcam] feature_maps : {feature_maps}')
bbox_tensors = []
for i, fm in enumerate(feature_maps):
bbox_tensors.append(decode(fm, NUM_CLASS, i))
model = tf.keras.Model(input_layer, bbox_tensors)
utils.load_weights(model, FLAGS.weights)
toc = time.perf_counter()
print(f'[INFO] Architecture built.')
print(f'[DEBUG][webcam] Execution took {(1000 * (toc - tic)):0.4f} ms')
vid = cv2.VideoCapture(0)
if save_path:
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))
print(f"[DEBUG][video] Video CODEC : {FLAGS.save_path.split('.')[1]}")
codec = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(FLAGS.save_path, codec, fps, (width, height))
while True:
return_value, frame = vid.read()
if return_value:
print(f'[DEBUG] Got video capture')
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame)
else:
raise ValueError("No image! Try with another video format")
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)
prev_time = time.perf_counter()
pred_bbox = model.predict(image_data)
print(f'[INFO][webcam] Finished initial predication on image')
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS, STRIDES,
XYSCALE)
bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size,
score_thresh)
bboxes = utils.nms(bboxes, iou_thresh, method='nms')
image = utils.draw_bbox(frame, bboxes)
curr_time = time.perf_counter()
exec_time = curr_time - prev_time
result = np.asarray(image)
info = "fdpms: %.2f ms" % (1000 * exec_time)
print(info)
cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imshow("result", result)
print(result.shape)
if save_path:
out.write(result)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
vid.release()
out.release()
def main(_argv):
INPUT_SIZE = FLAGS.size
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
CLASSES = utils.read_class_names(cfg.YOLO.CLASSES)
predicted_dir_path = './mAP/predicted'
ground_truth_dir_path = './mAP/ground-truth'
if os.path.exists(predicted_dir_path): shutil.rmtree(predicted_dir_path)
if os.path.exists(ground_truth_dir_path):
shutil.rmtree(ground_truth_dir_path)
if os.path.exists(cfg.TEST.DECTECTED_IMAGE_PATH):
shutil.rmtree(cfg.TEST.DECTECTED_IMAGE_PATH)
os.mkdir(predicted_dir_path)
os.mkdir(ground_truth_dir_path)
os.mkdir(cfg.TEST.DECTECTED_IMAGE_PATH)
times = []
classes = [
'Book', 'Bottle', 'Computer keyboard', 'Computer mouse', 'Laptop',
'Mobile phone', 'Backpack'
]
# Build Model
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([INPUT_SIZE, INPUT_SIZE, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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)
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)
num_lines = sum(1 for line in open(FLAGS.annotation_path))
with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file:
for num, line in enumerate(annotation_file):
annotation = line.strip().split()
image_path = annotation[0]
image_name = image_path.split('/')[-1]
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
bbox_data_gt = np.array(
[list(map(float, box.split(','))) for box in annotation[1:]])
if len(bbox_data_gt) == 0:
bboxes_gt = []
classes_gt = []
else:
bboxes_gt, classes_gt = bbox_data_gt[:, :4], bbox_data_gt[:, 4]
ground_truth_path = os.path.join(ground_truth_dir_path,
str(num) + '.txt')
current_class = ''
print('=> ground truth of %s:' % image_name)
num_bbox_gt = len(bboxes_gt)
with open(ground_truth_path, 'w') as f:
for i in range(num_bbox_gt):
# esto
class_name = CLASSES[classes_gt[i]]
# esto
if i == 0:
current_class = class_name
class_name = CLASSES[classes_gt[i]]
if class_name == current_class:
xmin, ymin, xmax, ymax = list(map(str, bboxes_gt[i]))
bbox_mess = ' '.join(
[class_name, xmin, ymin, xmax, ymax]) + '\n'
f.write(bbox_mess)
print('\t' + str(bbox_mess).strip())
print('=> predict result of %s:' % image_name)
predict_result_path = os.path.join(predicted_dir_path,
str(num) + '.txt')
# Predict Process
image_size = image.shape[:2]
image_data = utils.image_preprocess(np.copy(image),
[INPUT_SIZE, INPUT_SIZE])
image_data = image_data[np.newaxis, ...].astype(np.float32)
if FLAGS.framework == "tf":
startTime = time.time()
pred_bbox = model.predict(image_data)
times.append(time.time() - startTime)
else:
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 == 'yolov3':
pred_bbox = utils.postprocess_bbbox(pred_bbox, ANCHORS,
STRIDES)
elif FLAGS.model == 'yolov4':
XYSCALE = cfg.YOLO.XYSCALE
pred_bbox = utils.postprocess_bbbox(pred_bbox,
ANCHORS,
STRIDES,
XYSCALE=XYSCALE)
pred_bbox = tf.concat(pred_bbox, axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, image_size, INPUT_SIZE,
cfg.TEST.SCORE_THRESHOLD)
bboxes = utils.nms(bboxes, cfg.TEST.IOU_THRESHOLD, method='nms')
if cfg.TEST.DECTECTED_IMAGE_PATH is not None:
image = utils.draw_bbox(image, bboxes)
cv2.imwrite(cfg.TEST.DECTECTED_IMAGE_PATH + image_name, image)
with open(predict_result_path, 'w') as f:
for bbox in bboxes:
if (CLASSES[int(bbox[5])] in classes) and (current_class
== CLASSES[int(
bbox[5])]):
coor = np.array(bbox[:4], dtype=np.int32)
score = bbox[4]
class_ind = int(bbox[5])
class_name = CLASSES[class_ind]
score = '%.4f' % score
xmin, ymin, xmax, ymax = list(map(str, coor))
bbox_mess = ' '.join(
[class_name, score, xmin, ymin, xmax, ymax]) + '\n'
f.write(bbox_mess)
print('\t' + str(bbox_mess).strip())
print(num, num_lines)
print("Elapsed time: " + str(sum(times) / len(times)))
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):
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
input_size = FLAGS.size
video_path = FLAGS.video
print("Video from: ", video_path)
vid = cv2.VideoCapture(video_path)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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)
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()
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)
while True:
return_value, frame = vid.read()
if return_value:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame)
else:
cv2.destroyWindow("result")
raise ValueError("No image! Try with another video format")
while cv2.getWindowProperty('window-name', 0) >= 0:
keyCode = cv2.waitKey(50)
#break
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)
prev_time = time.time()
if FLAGS.framework == 'tf':
pred_bbox = model.predict(image_data)
else:
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':
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.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
image = utils.draw_bbox(frame, bboxes)
curr_time = time.time()
exec_time = curr_time - prev_time
result = np.asarray(image)
info = "time: %.2f ms" % (1000 * exec_time)
print(info)
cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
result = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imshow("result", result)
if cv2.waitKey(1) & 0xFF == ord('q'): break
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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
input_size = FLAGS.size
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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 = []
def main(_argv):
if FLAGS.tiny:
STRIDES = np.array(cfg.YOLO.STRIDES_TINY)
ANCHORS = utils.get_anchors(cfg.YOLO.ANCHORS_TINY, FLAGS.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)
NUM_CLASS = len(utils.read_class_names(cfg.YOLO.CLASSES))
XYSCALE = cfg.YOLO.XYSCALE
input_size = FLAGS.size
video_path = FLAGS.video
print("Video from: ", video_path)
vid = cv2.VideoCapture(video_path)
if FLAGS.framework == 'tf':
input_layer = tf.keras.layers.Input([input_size, input_size, 3])
if FLAGS.tiny:
feature_maps = YOLOv3_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)
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()
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)
# setup for output video
width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH) + 0.5)
height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT) + 0.5)
size = (width, height)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter('/content/output-vid.avi', fourcc, 20.0, size)
total_frames = int(vid.get(cv2.CAP_PROP_FRAME_COUNT))
print('Total Frames:', total_frames)
while True:
return_value, frame = vid.read()
n_frame = int(vid.get(cv2.CAP_PROP_POS_FRAMES))
if return_value:
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(frame)
else:
out.release()
if (total_frames - 1 != n_frame):
raise ValueError("No image! Try with another video format")
print("Finished processing video.")
break
frame_size = frame.shape[:2]
image_data = utils.image_preporcess(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)
else:
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':
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.25)
bboxes = utils.nms(bboxes, 0.213, method='nms')
image = utils.draw_bbox(frame, bboxes)
curr_time = time.time()
exec_time = curr_time - prev_time
result = np.asarray(image)
info = "time: %.2f ms" % (1000 * exec_time)
print(info)
# write modified frame to video
resultFrame = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
out.write(resultFrame)
# save modified frames
print("Frame:", n_frame)
cv2.imwrite("frame{}.jpg".format(n_frame), resultFrame)
if cv2.waitKey(1) & 0xFF == ord('q'): break
