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()
示例#3
0
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'))
示例#5
0
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()