def get_eval_model(self):
# Load model, or construct model and load weights.
num_anchors = len(self.anchors)
num_classes = len(self.class_names)
#YOLOv3 model has 9 anchors and 3 feature layers but
#Tiny YOLOv3 model has 6 anchors and 2 feature layers,
#so we can calculate feature layers number to get model type
num_feature_layers = num_anchors // 3
if self.model_type.startswith(
'scaled_yolo4_') or self.model_type.startswith('yolo5_'):
# Scaled-YOLOv4 & YOLOv5 entrance
eval_model, _ = get_yolo5_model(self.model_type,
num_feature_layers,
num_anchors,
num_classes,
input_shape=self.model_image_size +
(3, ),
model_pruning=self.model_pruning)
self.v5_decode = True
elif self.model_type.startswith('yolo3_') or self.model_type.startswith('yolo4_') or \
self.model_type.startswith('tiny_yolo3_') or self.model_type.startswith('tiny_yolo4_'):
# YOLOv3 & v4 entrance
eval_model, _ = get_yolo3_model(self.model_type,
num_feature_layers,
num_anchors,
num_classes,
input_shape=self.model_image_size +
(3, ),
model_pruning=self.model_pruning)
self.v5_decode = False
elif self.model_type.startswith(
'yolo2_') or self.model_type.startswith('tiny_yolo2_'):
# YOLOv2 entrance
eval_model, _ = get_yolo2_model(self.model_type,
num_anchors,
num_classes,
input_shape=self.model_image_size +
(3, ),
model_pruning=self.model_pruning)
self.v5_decode = False
else:
raise ValueError('Unsupported model type')
return eval_model
def _generate_model(self):
'''to generate the bounding boxes'''
weights_path = os.path.expanduser(self.weights_path)
assert weights_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)
#YOLOv3 model has 9 anchors and 3 feature layers but
#Tiny YOLOv3 model has 6 anchors and 2 feature layers,
#so we can calculate feature layers number to get model type
num_feature_layers = num_anchors // 3
try:
if self.model_type.startswith(
'scaled_yolo4_') or self.model_type.startswith('yolo5_'):
# Scaled-YOLOv4 & YOLOv5 entrance
yolo_model, _ = get_yolo5_model(
self.model_type,
num_feature_layers,
num_anchors,
num_classes,
input_shape=self.model_image_size + (3, ),
model_pruning=self.pruning_model)
elif self.model_type.startswith('yolo3_') or self.model_type.startswith('yolo4_') or \
self.model_type.startswith('tiny_yolo3_') or self.model_type.startswith('tiny_yolo4_'):
# YOLOv3 & v4 entrance
yolo_model, _ = get_yolo3_model(
self.model_type,
num_feature_layers,
num_anchors,
num_classes,
input_shape=self.model_image_size + (3, ),
model_pruning=self.pruning_model)
elif self.model_type.startswith(
'yolo2_') or self.model_type.startswith('tiny_yolo2_'):
# YOLOv2 entrance
yolo_model, _ = get_yolo2_model(
self.model_type,
num_anchors,
num_classes,
input_shape=self.model_image_size + (3, ),
model_pruning=self.pruning_model)
else:
raise ValueError('Unsupported model type')
yolo_model.load_weights(
weights_path) # make sure model, anchors and classes match
if self.pruning_model:
yolo_model = sparsity.strip_pruning(yolo_model)
yolo_model.summary()
except Exception as e:
print(repr(e))
assert yolo_model.layers[-1].output_shape[-1] == \
num_anchors/len(yolo_model.output) * (num_classes + 5), \
'Mismatch between model and given anchor and class sizes'
print('{} model, anchors, and classes loaded.'.format(weights_path))
if self.gpu_num >= 2:
yolo_model = multi_gpu_model(yolo_model, gpus=self.gpu_num)
return yolo_model