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