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