def main(): # 加载数据 image_wh = (416, 416) # image_wh = (608, 608) anchors = LoadAnchors(anchorsFile) classes_name, classes_num = LoadClasses(classesFile) # 构建模型 model = YoloV3Model(classes_num=classes_num, anchors=anchors, image_wh=image_wh) # 编译模型 print('编译模型') model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4)) # 加载模型 _ = model(tf.ones((1, image_wh[1], image_wh[0], 3)), training=False) if os.path.exists(modelPath): last_model_path = tf.train.latest_checkpoint(modelPath) model.load_weights(last_model_path).expect_partial() print('加载模型:{}'.format(last_model_path)) # model.summary() # 保存模型 model.save_weights(teacherModelPath) print('保存模型:{}'.format(teacherModelPath))
def main(): anchors = LoadAnchors('./data/coco_anchors.txt') data_generator = DataGenerator(image_path='Z:/Labels/coco2017/train2017', label_path='.\\data\\coco_train2017_labels.txt', classes_path='.\\data\\coco_classes.txt', batch_size=3, anchors=anchors, image_wh=(416, 416), label_mean=True, flip=True) dataset = data_generator.GetDataSet()
import numpy as np import tensorflow as tf import shutil import h5py import argparse import sys import os sys.path.append(os.getcwd()) from ai_api.ai_models.unsupervised_learning.model import YoloV3ModelBase from ai_api.ai_models.utils.load_object_detection_data import LoadClasses, LoadAnchors parser = argparse.ArgumentParser() parser.add_argument( '--output_path', default='./data/unsupervised_learning_weights/tf2_weights/tf2_weights.ckpt' ) parser.add_argument('--classes_num', default=80, type=int, help='类别数量') args = parser.parse_args() tf.random.set_seed(991) anchors = LoadAnchors('./data/coco_anchors.txt') model = YoloV3ModelBase(classes_num=args.classes_num, anchors_num=anchors.shape[1], backbone_weights='imagenet') y = model(tf.zeros([1, 416, 416, 3], tf.float32), training=True) model.summary() model.save_weights(args.output_path)
def train(): '''训练''' # 设置GPU显存自适应 gpus = tf.config.experimental.list_physical_devices(device_type='GPU') cpus = tf.config.experimental.list_physical_devices(device_type='CPU') print(gpus, cpus) for gpu in gpus: tf.config.experimental.set_memory_growth(gpu, True) # if len(gpus) > 1: # tf.config.experimental.set_visible_devices(gpus[1], 'GPU') # elif len(cpus) > 0: # tf.config.experimental.set_visible_devices(cpus[0], 'CPU') # 加载数据 anchors = LoadAnchors(anchorsFile) data_generator_train = TeacherDataGenerator(image_path=trainData, classes_path=classesFile, batch_size=batchSize, anchors=anchors, image_wh=(416, 416), image_random=True, flip=True) data_set_train = data_generator_train.GetDataSet() data_generator_val = DataGenerator(image_path=valData, label_path=valLabels, classes_path=classesFile, batch_size=1, anchors=anchors, image_wh=(416, 416), label_mean=False, image_random=False, flip=False) data_set_val = data_generator_val.GetDataSet() # 构建模型 model = YoloV3Model(classes_num=data_generator_train.classes_num, anchors=anchors, image_wh=(416, 416), backbone_weights=None) # 编译模型 print('编译模型') model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4)) # model.compile(optimizer=RAdam(lr=1e-4)) # 日志 log_dir = './data/unsupervised_learning_weights/' model_path = log_dir + 'train_weights/' old_model_path = log_dir + 'tf2_weights/' _ = model(tf.ones((1, 416, 416, 3))) is_old_model = False if os.path.exists(model_path): last_model_path = tf.train.latest_checkpoint(model_path) model.load_weights(last_model_path) print('加载模型:{}'.format(last_model_path)) elif os.path.exists(old_model_path): last_model_path = tf.train.latest_checkpoint(old_model_path) model.load_weights(last_model_path) print('加载模型:{}'.format(last_model_path)) is_old_model = True model.summary() # 训练回调方法 # logging = tf.keras.callbacks.TensorBoard(log_dir=log_dir) checkpoint = tf.keras.callbacks.ModelCheckpoint(os.path.join( model_path, 'ep{epoch:04d}-loss{loss:.3f}-val_mAP{val_mAP:.3f}.ckpt'), monitor='loss', mode='min', save_weights_only=True, save_best_only=False, verbose=1) reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(monitor='loss', factor=0.1, patience=3, verbose=1) early_stopping = tf.keras.callbacks.EarlyStopping(monitor='loss', min_delta=0, patience=10, verbose=1) if is_old_model: print('旧权重,开始预训练') model.FreeLayer( ['darknet_conv2d_6', 'darknet_conv2d_14', 'darknet_conv2d_22']) model.fit(data_set_train, steps_per_epoch=1000, epochs=1, initial_epoch=0) model.save_weights(os.path.join(model_path, 'start_weights.ckpt')) print('Train on {} samples, val on {} samples, with batch size {}.'.format( data_generator_train.image_num, data_generator_val.labels_num, batchSize)) print('开始训练') steps_per_epoch = 5000 model.FreeLayer(['']) model.fit( data_set_train, # steps_per_epoch=max(1, data_generator_train.labels_num//batchSize), steps_per_epoch=steps_per_epoch, validation_data=data_set_val, validation_steps=max(1, data_generator_val.labels_num), # validation_steps=10, epochs=300, # initial_epoch=(model.global_step.numpy()//steps_per_epoch), initial_epoch=(model.optimizer.iterations.numpy() // steps_per_epoch), callbacks=[reduce_lr, early_stopping, checkpoint]) model.save_weights(os.path.join(model_path, 'last_weights.ckpt'))
def test(): '''训练''' # 加载数据 image_wh = (416, 416) # image_wh = (608, 608) anchors = LoadAnchors(anchorsFile) classes_name, classes_num = LoadClasses(classesFile) # 构建模型 model = YoloV3Model(classes_num=classes_num, anchors=anchors, image_wh=image_wh) # 编译模型 print('编译模型') model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4)) # 加载模型 _ = model(tf.ones((1, image_wh[1], image_wh[0], 3)), training=False) if os.path.exists(modelPath): last_model_path = tf.train.latest_checkpoint(modelPath) model.load_weights(last_model_path).expect_partial() print('加载模型:{}'.format(last_model_path)) # model.summary() img_old = ImageHelper.fileToOpencvImage(imageFile) # 缩放图片 img, _, padding = ImageHelper.opencvProportionalResize(img_old, image_wh, bg_color=(0, 0, 0)) # print('imgType:', type(img)) width, height = ImageHelper.opencvGetImageSize(img_old) image_size_old = np.int32([width, height]) print('imgSize:', width, height) # 最后输出图片 predict_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # 调整参数范围 predict_img = predict_img.astype(np.float32) predict_img = predict_img / 255 # 增加一个维度 predict_img = np.expand_dims(predict_img, 0) y_boxes, y_classes_id, y_scores, y_classes, y_confidence = model.Predict( predict_img, confidence_thresh=0.5, scores_thresh=0.2, iou_thresh=0.5) # 结果转物体框列表 y_boxes = y_boxes.numpy() y_classes_id = y_classes_id.numpy() y_scores = y_scores.numpy() y_classes = y_classes.numpy() y_confidence = y_confidence.numpy() y_boxes[:, [0, 2]] = (y_boxes[:, [0, 2]] * image_wh[0] - padding[2]) / ( image_wh[0] - padding[2] - padding[3]) * image_size_old[0] y_boxes[:, [1, 3]] = (y_boxes[:, [1, 3]] * image_wh[1] - padding[0]) / ( image_wh[1] - padding[0] - padding[1]) * image_size_old[1] # 截取框超出图片部分 y_boxes[:, 0][y_boxes[:, 0] < 0] = 0 y_boxes[:, 1][y_boxes[:, 1] < 0] = 0 y_boxes[:, 2][y_boxes[:, 2] > image_size_old[0]] = image_size_old[0] y_boxes[:, 3][y_boxes[:, 3] > image_size_old[1]] = image_size_old[1] # 去掉无效框 y_mask = np.logical_and(y_boxes[:, 2] - y_boxes[:, 0] > 2, y_boxes[:, 3] - y_boxes[:, 1] > 2) y_boxes = y_boxes[y_mask] y_classes_id = y_classes_id[y_mask] y_scores = y_scores[y_mask] y_classes = y_classes[y_mask] y_confidence = y_confidence[y_mask] y_boxes = y_boxes.astype(np.int32) # print('y_boxes:', y_boxes.shape) # print('y_classes_id:', y_classes_id.shape) # print('y_scores:', y_scores.shape) # print('y_classes:', y_classes.shape) # print('y_confidence:', y_confidence.shape) colors = [] for cr in range(0, 256, 32): for cg in range(0, 256, 32): for cb in range(0, 256, 32): colors.append([cb, cg, cr]) colors = np.array(colors) np.random.seed(11) np.random.shuffle(colors) np.random.seed(None) result_img = img_old.copy() for i in range(y_boxes.shape[0]): # print('y_boxes:', y_boxes[i,:]) # print('y_classes_id:', y_classes_id[i]) # print('y_classes_name:', classes_name[y_classes_id[i]]) # print('y_scores:', y_scores[i]) # print('y_classes:', y_classes[i]) # print('y_confidence:', y_confidence[i]) cv2.rectangle(result_img, tuple(y_boxes[i, 0:2]), tuple(y_boxes[i, 2:4]), colors[y_classes_id[i] % len(colors)].tolist(), thickness=1) cv2.putText(result_img, classes_name[y_classes_id[i]], tuple(y_boxes[i, 0:2]), cv2.FONT_HERSHEY_COMPLEX, 0.6, (0, 100, 0), 1) cv2.putText(result_img, str(y_scores[i]), tuple(y_boxes[i, 0:2] + (0, 20)), cv2.FONT_HERSHEY_COMPLEX, 0.6, (0, 0, 100), 1) ImageHelper.showOpencvImage(result_img)
import numpy as np import math import time import os import tensorflow as tf from ai_api.ai_models.yolo_v4.model import YoloV4Model import ai_api.ai_models.utils.image_helper as ImageHelper from ai_api.ai_models.utils.load_object_detection_data import LoadClasses, LoadAnchors modelPath = './data/yolo_v3_weights/tf2_weights/' classesFile = './data/coco_classes.txt' anchorsFile = './data/coco_anchors.txt' # 加载数据 image_wh = (416, 416) # image_wh = (608, 608) anchors = LoadAnchors(anchorsFile) classes_name, classes_num = LoadClasses(classesFile) # 构建模型 model = YoloV4Model(classes_num=classes_num, anchors=anchors, image_wh=image_wh) # 编译模型 print('编译模型') model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4)) # 加载模型 _ = model(tf.ones((1, image_wh[1], image_wh[0], 3)), training=False) if os.path.exists(modelPath): last_model_path = tf.train.latest_checkpoint(modelPath) model.load_weights(last_model_path).expect_partial()