def main(_argv):
print('Config File From:', FLAGS.config)
cfg = decode_cfg(FLAGS.config)
model_type = cfg['yolo']['type']
if model_type == 'yolov3':
from core.model.one_stage.yolov3 import YOLOv3 as Model
from core.model.one_stage.yolov3 import YOLOLoss as Loss
num = 186
epochs = 270
elif model_type == 'yolov3_tiny':
from core.model.one_stage.yolov3 import YOLOv3_Tiny as Model
from core.model.one_stage.yolov3 import YOLOLoss as Loss
num = 29
epochs = 30
elif model_type == 'yolov4':
from core.model.one_stage.yolov4 import YOLOv4 as Model
from core.model.one_stage.yolov4 import YOLOLoss as Loss
num = 251
epochs = 270
elif model_type == 'yolov4_tiny':
from core.model.one_stage.yolov4 import YOLOv4_Tiny as Model
from core.model.one_stage.yolov4 import YOLOLoss as Loss
num = 29
epochs = 30
else:
raise NotImplementedError()
model, eval_model = Model(cfg)
model.summary()
train_dataset = Dataset(cfg)
init_weight = cfg["train"]["init_weight_path"]
anchors = cfg['yolo']['anchors']
mask = cfg['yolo']['mask']
strides = cfg['yolo']['strides']
ignore_threshold = cfg['train']['ignore_threshold']
loss_type = cfg['train']['loss_type']
if init_weight:
load_weights(model, init_weight)
else:
print("Training from scratch")
num = 0
loss = [
Loss(anchors[mask[i]], strides[i], train_dataset.num_classes,
ignore_threshold, loss_type) for i in range(len(mask))
]
ckpt_path = os.path.join(cfg["train"]["save_weight_path"], 'tmp',
cfg["train"]["label"],
time.strftime("%Y%m%d%H%M", time.localtime()))
warmup_epochs = 3
warmup_callback = [
WarmUpScheduler(learning_rate=1e-3,
warmup_step=warmup_epochs * len(train_dataset),
verbose=1)
]
eval_callback = [
COCOEvalCheckpoint(save_path=os.path.join(ckpt_path,
"mAP-{mAP:.4f}.h5"),
eval_model=eval_model,
model_cfg=cfg,
eval_n_samples=None,
eval_per_batch=32000,
verbose=1)
]
lr_callback = [
CosineAnnealingScheduler(learning_rate=1e-3,
eta_min=1e-6,
T_max=epochs * len(train_dataset),
verbose=1)
]
if not os.path.isdir(ckpt_path):
os.makedirs(ckpt_path)
os.makedirs(os.path.join(ckpt_path, 'train', 'plugins', 'profile'))
# opt = optimizers.Adam(lr=0.)
opt = Accumulative(optimizers.Adam(lr=0.), 32)
# warm-up
for i in range(num):
model.layers[i].trainable = False
print(model.layers[i].name)
print('Freeze the first {} layers of total {} layers.'.format(
num, len(model.layers)))
model.compile(loss=loss, optimizer=opt, run_eagerly=False)
model.fit(train_dataset,
steps_per_epoch=len(train_dataset),
epochs=warmup_epochs,
callbacks=warmup_callback)
for i in range(len(model.layers)):
model.layers[i].trainable = True
print('Unfreeze all layers.')
model.compile(loss=loss, optimizer=opt, run_eagerly=False)
model.fit(train_dataset,
steps_per_epoch=len(train_dataset),
epochs=epochs,
callbacks=eval_callback + lr_callback)
# reset sample rate
model.compile(loss=loss,
optimizer=optimizers.Adam(lr=1e-7),
run_eagerly=False)
model.fit(train_dataset,
steps_per_epoch=len(train_dataset),
epochs=10,
callbacks=eval_callback)
def main(_argv):
cfg = decode_cfg("cfgs/voc_yolov4_tiny.yaml")
model, eval_model = YOLOv4_Tiny(cfg)
model.summary()
train_dataset = Dataset(cfg)
init_weight = cfg["train"]["init_weight_path"]
anchors = cfg['yolo']['anchors']
mask = cfg['yolo']['mask']
strides = cfg['yolo']['strides']
ignore_threshold = cfg['train']['ignore_threshold']
loss_type = cfg['train']['loss_type']
if init_weight:
load_weights(model, init_weight)
else:
print("Training from scratch")
loss = [
YOLOLoss(anchors[mask[i]], strides[i], train_dataset.num_classes,
ignore_threshold, loss_type) for i in range(len(mask))
]
ckpt_path = os.path.join(cfg["train"]["save_weight_path"], 'tmp',
cfg["train"]["label"],
time.strftime("%Y%m%d%H%M", time.localtime()))
if not os.path.isdir(ckpt_path):
os.makedirs(ckpt_path)
os.makedirs(
os.path.join(ckpt_path, 'log', 'train', 'plugins', 'profile'))
_cfg = copy.deepcopy(cfg)
_cfg['test']['anno_path'] = "./data/pascal_voc/voc2007_val.txt"
callback = [
COCOEvalCheckpoint(save_path=os.path.join(ckpt_path,
"mAP-{mAP:.4f}.h5"),
eval_model=eval_model,
model_cfg=cfg,
sample_rate=5,
verbose=1),
COCOEvalCheckpoint(save_path=None,
eval_model=eval_model,
model_cfg=_cfg,
sample_rate=5,
verbose=1)
]
num = 29
for i in range(num):
model.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(
num, len(model.layers)))
model.compile(loss=loss,
optimizer=optimizers.Adam(lr=1e-4),
run_eagerly=False)
model.fit(train_dataset,
steps_per_epoch=len(train_dataset),
epochs=30,
callbacks=callback)
for i in range(len(model.layers)):
model.layers[i].trainable = True
model.compile(loss=loss,
optimizer=optimizers.Adam(lr=1e-5),
run_eagerly=False)
model.fit(train_dataset,
steps_per_epoch=len(train_dataset),
epochs=50,
callbacks=callback)
callback = [
COCOEvalCheckpoint(save_path=os.path.join(ckpt_path,
"mAP-{mAP:.4f}.h5"),
eval_model=eval_model,
model_cfg=cfg,
sample_rate=1,
verbose=1),
]
model.compile(loss=loss,
optimizer=optimizers.Adam(lr=1e-6),
run_eagerly=False)
model.fit(train_dataset,
steps_per_epoch=len(train_dataset),
epochs=10,
callbacks=callback)
def main(_argv):
# read config
print('Config File From:', FLAGS.config)
print('Media From:', FLAGS.media)
print('Use GPU:', FLAGS.gpu)
if not FLAGS.gpu:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
cfg = decode_cfg(FLAGS.config)
model_type = cfg['yolo']['type']
if model_type == 'yolov3':
from core.model.one_stage.yolov3 import YOLOv3 as Model
elif model_type == 'yolov3_tiny':
from core.model.one_stage.yolov3 import YOLOv3_Tiny as Model
elif model_type == 'yolov4':
from core.model.one_stage.yolov4 import YOLOv4 as Model
elif model_type == 'yolov4_tiny':
from core.model.one_stage.yolov4 import YOLOv4_Tiny as Model
elif model_type == 'yolox':
from core.model.one_stage.custom import YOLOX as Model
else:
raise NotImplementedError()
_, model = Model(cfg)
model.summary()
init_weight_path = cfg['test']['init_weight_path']
if init_weight_path:
print('Load Weights File From:', init_weight_path)
load_weights(model, init_weight_path)
else:
raise SystemExit('init_weight_path is Empty !')
# assign colors for difference labels
shader = Shader(cfg['yolo']['num_classes'])
names = cfg['yolo']['names']
image_size = cfg['test']['image_size'][0]
#model.save('E:/dm/repo/yolox/ckpts/tmp/voc_yolov4_tiny_SM_DM_CIoU_FL/yolov4_tiny_best/yolov4_tiny.h5')
# full_model = tf.function(lambda Input: model(Input))
# full_model = full_model.get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
#
# # Get frozen ConcreteFunction
# frozen_func = convert_variables_to_constants_v2(full_model)
# frozen_func.graph.as_graph_def()
#
# layers = [op.name for op in frozen_func.graph.get_operations()]
# print("-" * 50)
# print("Frozen model layers: ")
# for layer in layers:
# print(layer)
#
# print("-" * 50)
# print("Frozen model inputs: ")
# print(frozen_func.inputs)
# print("Frozen model outputs: ")
# print(frozen_func.outputs)
#
# # Save frozen graph from frozen ConcreteFunction to hard drive
# tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
# logdir="./frozen_models",
# name="yolov4_tiny_tf.pb",
# as_text=False)
def inference(image):
h, w = image.shape[:2]
image = preprocess_image(image,
(image_size, image_size)).astype(np.float32)
images = np.expand_dims(image, axis=0)
tic = time.time()
bboxes, scores, classes, valid_detections = model.predict(images)
toc = time.time()
bboxes = bboxes[0][:valid_detections[0]]
scores = scores[0][:valid_detections[0]]
classes = classes[0][:valid_detections[0]]
# bboxes *= image_size
_, bboxes = postprocess_image(image, (w, h), bboxes)
return (toc - tic) * 1000, bboxes, scores, classes
if FLAGS.media.startswith('rtsp') or FLAGS.media.isdigit(
) or FLAGS.media.endswith('.mp4') or FLAGS.media.endswith('.avi'):
from collections import deque
d = deque(maxlen=10)
media = read_video(FLAGS.media)
while True:
ret, image = media.read()
if not ret: break
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
ms, bboxes, scores, classes = inference(image)
image = draw_bboxes(image, bboxes, scores, classes, names, shader)
d.append(ms)
mms = np.mean(d)
print('Inference Time:', mms, 'ms')
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
image = cv2.putText(image, "{:.2f} ms".format(mms), (0, 30),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 255, 0),
2)
cv2.imshow('Image', image)
if cv2.waitKey(33) == ord('q'):
break
media.release()
elif FLAGS.media.endswith('.jpg') or FLAGS.media.endswith('.png'):
image = read_image(FLAGS.media)
ms, bboxes, scores, classes = inference(image)
image = draw_bboxes(image, bboxes, scores, classes, names, shader)
print('Inference Time:', ms, 'ms')
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imshow('Image', image)
cv2.waitKey()
#!/usr/bin/env python
# -*- coding: utf-8 -*-
from core.utils import decode_cfg, load_weights
from core.callbacks import COCOEvalCheckpoint, VOCEvalCheckpoint
if __name__ == '__main__':
cfg = decode_cfg('cfgs/coco_yolov3_tiny.yaml')
model_type = cfg['yolo']['type']
if model_type == 'yolov3':
from core.model.one_stage.yolov3 import YOLOv3 as Model
elif model_type == 'yolov3_tiny':
from core.model.one_stage.yolov3 import YOLOv3_Tiny as Model
elif model_type == 'yolov4':
from core.model.one_stage.yolov4 import YOLOv4 as Model
elif model_type == 'yolov4_tiny':
from core.model.one_stage.yolov4 import YOLOv4_Tiny as Model
else:
raise NotImplementedError()
model, eval_model = Model(cfg)
model.summary()
init_weight = cfg["train"]["init_weight_path"]
load_weights(model, init_weight)
def main(_argv):
# read config
print('Config File From:', FLAGS.config)
print('Media From:', FLAGS.media)
print('Use GPU:', FLAGS.gpu)
if not FLAGS.gpu:
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
cfg = decode_cfg(FLAGS.config)
model_type = cfg['yolo']['type']
if model_type == 'yolov3':
from core.model.one_stage.yolov3 import YOLOv3 as Model
elif model_type == 'yolov3_tiny':
from core.model.one_stage.yolov3 import YOLOv3_Tiny as Model
elif model_type == 'yolov4':
from core.model.one_stage.yolov4 import YOLOv4 as Model
elif model_type == 'yolov4_tiny':
from core.model.one_stage.yolov4 import YOLOv4_Tiny as Model
else:
raise NotImplementedError()
_, model = Model(cfg)
model.summary()
init_weight_path = cfg['test']['init_weight_path']
if init_weight_path:
print('Load Weights File From:', init_weight_path)
load_weights(model, init_weight_path)
else:
raise SystemExit('init_weight_path is Empty !')
# assign colors for difference labels
shader = Shader(cfg['yolo']['num_classes'])
names = cfg['yolo']['names']
image_size = cfg['test']['image_size'][0]
def inference(image):
h, w = image.shape[:2]
image = preprocess_image(image,
(image_size, image_size)).astype(np.float32)
images = np.expand_dims(image, axis=0)
tic = time.time()
bboxes, scores, classes, valid_detections = model.predict(images)
toc = time.time()
bboxes = bboxes[0][:valid_detections[0]]
scores = scores[0][:valid_detections[0]]
classes = classes[0][:valid_detections[0]]
# bboxes *= image_size
_, bboxes = postprocess_image(image, (w, h), bboxes)
return (toc - tic) * 1000, bboxes, scores, classes
if FLAGS.media.startswith('rtsp') or FLAGS.media.isdigit(
) or FLAGS.media.endswith('.mp4') or FLAGS.media.endswith('.avi'):
from collections import deque
d = deque(maxlen=10)
media = read_video(FLAGS.media)
while True:
ret, image = media.read()
if not ret: break
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
ms, bboxes, scores, classes = inference(image)
image = draw_bboxes(image, bboxes, scores, classes, names, shader)
d.append(ms)
mms = np.mean(d)
print('Inference Time:', mms, 'ms')
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
image = cv2.putText(image, "{:.2f} ms".format(mms), (0, 30),
cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 255, 0),
2)
cv2.imshow('Image', image)
if cv2.waitKey(33) == ord('q'):
break
media.release()
elif FLAGS.media.endswith('.jpg') or FLAGS.media.endswith('.png'):
image = read_image(FLAGS.media)
ms, bboxes, scores, classes = inference(image)
image = draw_bboxes(image, bboxes, scores, classes, names, shader)
print('Inference Time:', ms, 'ms')
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
cv2.imshow('Image', image)
cv2.waitKey()
