def generate(self):
"""未运算, 返回model结果的tensor"""
weight_path = os.path.expanduser(self.weight_path)
assert weight_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
# Load model, or construct model and load weights.
num_anchors = len(self.anchors)
num_classes = len(self.class_names)
h, w = self.model_image_size
assert os.path.exists(self.weight_path), 'weight文件不存在'
self.yolo_model = yolo_body(Input(shape=(h, w, 3)), num_anchors // 3,
num_classes)
self.yolo_model.load_weights(
self.weight_path) # make sure model, anchors and classes match
print('{} model, anchors, and classes loaded.'.format(weight_path))
# Generate output tensor targets for filtered bounding boxes.
# placeholder: 实例化一个占位符张量并返回它
# input_image_shape 只输入wh
self.input_image_shape = K.placeholder(shape=(2, ))
if self.gpu_num >= 2:
self.yolo_model = multi_gpu_model(self.yolo_model,
gpus=self.gpu_num)
boxes, scores, classes = yolo_eval(self.yolo_model.output,
self.anchors,
len(self.class_names),
self.input_image_shape,
score_threshold=self.score)
# boxes: xywh
return boxes, scores, classes
def create_model(input_shape,
anchors,
num_class,
weights_path,
freeze_body=0,
load_pretrained=False):
"""
创建一个model
:param freeze_body: 冻结部分参数
:param input_shape: 输入尺寸, hw
:param anchors: 所有的anchors
:param num_class: class数量
:param weights_path: 已训练的权重的数量
:param load_pretrained: 是否加载权重
:return:
"""
K.clear_session()
image_input = Input(shape=(416, 416, 3))
h, w = input_shape
num_anchors = len(anchors)
out_dict = {0: 32, 1: 16, 2: 8}
# y_true的输出
y_true = [
Input(shape=(h // out_dict[m], w // out_dict[m], num_anchors // 3,
num_class + 5)) for m in range(3)
]
model_body = yolo_body(image_input, num_anchors // 3, num_class)
# 加载已训练的weight
if load_pretrained:
import os
assert os.path.exists(weights_path), 'weight文件不存在'
model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
print('Load weights {}.'.format(weights_path))
if freeze_body in [1, 2]:
# Freeze darknet53 body or freeze all but 3 output layers.
num = (185, len(model_body.layers) - 3)[freeze_body - 1]
for i in range(num):
model_body.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(
num, len(model_body.layers)))
loss_argument = {
'anchors': anchors,
'num_classes': num_class,
'ignore_thresh': 0.5
}
# yolo_loss封装为Layer对象(用来间接到yolo_body上, 使model的输出即为loss)
# 因为这个model只用来训练
model_loss = Lambda(yolo_loss,
output_shape=(1, ),
name='yolo_loss',
arguments=loss_argument)([*model_body.output, *y_true])
model = Model([model_body.input, *y_true], model_loss)
return model
def generate(self):
model_path = os.path.expanduser(self.model_path)
assert model_path.endswith(
'.h5'), 'Keras model or weights must be a .h5 file.'
# Load model, or construct model and load weights.
num_anchors = len(self.anchors)
num_classes = len(self.class_names)
is_tiny_version = num_anchors == 6 # default setting
try:
self.yolo_model = load_model(model_path, compile=False)
print('Model loaded')
except:
self.yolo_model = tiny_yolo_body(Input(shape=(None,None,3)), num_anchors//2, num_classes) \
if is_tiny_version else yolo_body(Input(shape=(None,None,3)), num_anchors//3, num_classes)
self.yolo_model.load_weights(
self.model_path) # make sure model, anchors and classes match
print('自己写模型')
else:
assert self.yolo_model.layers[-1].output_shape[-1] == \
num_anchors/len(self.yolo_model.output) * (num_classes + 5), \
'Mismatch between model and given anchor and class sizes'
print('{} model, anchors, and classes loaded.'.format(model_path))
# Generate colors for drawing bounding boxes.
hsv_tuples = [(x / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
self.colors))
np.random.seed(10101) # Fixed seed for consistent colors across runs.
np.random.shuffle(
self.colors) # Shuffle colors to decorrelate adjacent classes.
np.random.seed(None) # Reset seed to default.
# Generate output tensor targets for filtered bounding boxes.
self.input_image_shape = K.placeholder(shape=(2, ))
if self.gpu_num >= 2:
self.yolo_model = multi_gpu_model(self.yolo_model,
gpus=self.gpu_num)
boxes, scores, classes = yolo_eval(self.yolo_model.output,
self.anchors,
len(self.class_names),
self.input_image_shape,
score_threshold=self.score,
iou_threshold=self.iou)
return boxes, scores, classes
def _main():
# TODO: 定义路径
model_path = "logs/yolov3-model-1.h5"
assert model_path.endswith('.h5'), 'Keras model must be a .h5 file.'
anchors_path = "model/yolo_anchors.txt"
classes_path = "model/voc_classes.txt"
test_path = "images/test1.jpg"
output_path = "images/test1_out.jpg"
anchors = get_anchors(anchors_path)
class_names = get_classes(classes_path)
num_classes = len(class_names)
num_anchors = len(anchors)
# 构建模型
image_input = Input(shape=(None, None, 3))
model = yolo_body(image_input, num_anchors // 3, num_classes)
model.load_weights(model_path)
model.summary()
# 给每一个类定义一个颜色
hsv_tuples = [(x / num_classes, 1., 1.) for x in range(num_classes)]
colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
colors = list(
map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
colors))
random.seed(10101) # Fixed seed for consistent colors across runs.
random.shuffle(colors) # Shuffle colors to decorrelate adjacent classes.
random.seed(None) # Reset seed to default.
# TODO: 加载图片
image = Image.open(test_path)
print(image.size)
# 按照原尺寸缩放图像，空余的地方留灰
boxed_image = letterbox_image(image, (416, 416))
image_data = np.array(boxed_image, dtype='float32')
image_data /= 255.
image_data = np.expand_dims(image_data, 0) # Add batch dimension.
# 推理
y = model.predict(image_data, batch_size=1)
boxes_, scores_, classes_ = yolo_eval(y, anchors, num_classes,
(image.size[1], image.size[0]))
image = cv2.imread(test_path)
for i, box in enumerate(boxes_):
cv2.rectangle(image, (int(box[0]), int(box[1])),
(int(box[2]), int(box[3])), colors[classes_[i]])
cv2.putText(image, class_names[classes_[i]],
(int(box[0]), int(box[1])), 1, 1, colors[classes_[i]], 1)
cv2.imshow('image', image)
cv2.imwrite(output_path, image)
cv2.waitKey(0)
def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
weights_path='model/yolo_weights.h5'):
"""
create the training model
:param input_shape: 输入图片尺寸
:param anchors: anchor 列表
:param num_classes: classes 数量
:param load_pretrained: 是否加载预训练模型
:param freeze_body: 是否冻结部分层次
:param weights_path: 预训练模型路径
:return:
"""
K.clear_session()
image_input = Input(shape=(None, None, 3))
h, w = input_shape
num_anchors = len(anchors)
# [<tf.Tensor 'input_2:0' shape=(None, 13, 13, 3, 25) dtype=float32>,
# <tf.Tensor 'input_3:0' shape=(None, 26, 26, 3, 25) dtype=float32>,
# <tf.Tensor 'input_4:0' shape=(None, 52, 52, 3, 25) dtype=float32>]
y_true = [
Input(shape=(h // {0: 32, 1: 16, 2: 8}[l], w // {0: 32, 1: 16, 2: 8}[l], num_anchors // 3, num_classes + 5))
for
l in range(3)]
model_body = yolo_body(image_input, num_anchors // 3, num_classes)
print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))
# 加载域训练权重
if load_pretrained:
model_body.load_weights(weights_path, by_name=True, skip_mismatch=True)
print('Load weights {}.'.format(weights_path))
if freeze_body in [1, 2]:
# Freeze darknet53 body or freeze all but 3 output layers.
num = (185, len(model_body.layers) - 3)[freeze_body - 1]
for i in range(num):
model_body.layers[i].trainable = False
print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers)))
# todo 自定义loss
model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss',
arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})(
[*model_body.output, *y_true])
model = Model([model_body.input, *y_true], model_loss)
return model