def load_model(self, path):
self.yolo_model = models.load_model(path, compile=False)
yolov4_output = yolov4_head(self.yolo_model.output, self.num_classes,
self.anchors, self.xyscale)
self.inference_model = models.Model(
self.yolo_model.input,
nms(yolov4_output, self.img_size, self.num_classes
)) # [boxes, scores, classes, valid_detections]
def build_model(self, load_pretrained=True):
# core yolo model
input_layer = layers.Input(self.img_size)
yolov4_output = yolov4_neck(input_layer, self.num_classes)
self.yolo_model = models.Model(input_layer, yolov4_output)
# Build training model
y_true = [
layers.Input(name='input_2',
shape=(52, 52, 3,
(self.num_classes + 5))), # label small boxes
layers.Input(name='input_3',
shape=(26, 26, 3,
(self.num_classes + 5))), # label medium boxes
layers.Input(name='input_4',
shape=(13, 13, 3,
(self.num_classes + 5))), # label large boxes
layers.Input(name='input_5',
shape=(self.max_boxes, 4)), # true bboxes
]
loss_list = tf.keras.layers.Lambda(
yolo_loss,
name='yolo_loss',
arguments={
'num_classes': self.num_classes,
'iou_loss_thresh': self.iou_loss_thresh,
'anchors': self.anchors
})([*self.yolo_model.output, *y_true])
self.training_model = models.Model([self.yolo_model.input, *y_true],
loss_list)
# Build inference model
yolov4_output = yolov4_head(yolov4_output, self.num_classes,
self.anchors, self.xyscale)
# output: [boxes, scores, classes, valid_detections]
self.inference_model = models.Model(
input_layer,
nms(yolov4_output,
self.img_size,
self.num_classes,
iou_threshold=self.config['iou_threshold'],
score_threshold=self.config['score_threshold']))
if load_pretrained and self.weight_path and self.weight_path.endswith(
'.weights'):
if self.weight_path.endswith('.weights'):
load_weights(self.yolo_model, self.weight_path)
print(f'load from {self.weight_path}')
elif self.weight_path.endswith('.h5'):
self.training_model.load_weights(self.weight_path)
print(f'load from {self.weight_path}')
self.training_model.compile(optimizer=optimizers.Adam(lr=1e-3),
loss={
'yolo_loss':
lambda y_true, y_pred: y_pred
})
def predict_nonms(self,
img_path,
iou_threshold=0.413,
score_threshold=0.1):
raw_img = cv2.imread(img_path)
print('img shape: ', raw_img.shape)
img = self.preprocess_img(raw_img)
imgs = np.expand_dims(img, axis=0)
yolov4_output = self.yolo_model.predict(imgs)
output = yolov4_head(yolov4_output, self.num_classes, self.anchors,
self.xyscale)
pred_output = nms(output, self.img_size, self.num_classes,
iou_threshold, score_threshold)
pred_output = [p.numpy() for p in pred_output]
detections = get_detection_data(img=raw_img,
model_outputs=pred_output,
class_names=self.class_names)
draw_bbox(raw_img,
detections,
cmap=self.class_color,
random_color=True)
return detections