create the training model
(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2,
weights_path='model_data/yolo_weights.h5')
| 103 | |
| 104 | |
| 105 | def create_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2, |
| 106 | weights_path='model_data/yolo_weights.h5'): |
| 107 | '''create the training model''' |
| 108 | K.clear_session() # get a new session |
| 109 | image_input = Input(shape=(None, None, 3)) |
| 110 | h, w = input_shape |
| 111 | num_anchors = len(anchors) |
| 112 | |
| 113 | y_true = [Input(shape=(h//{0:32, 1:16, 2:8}[l], w//{0:32, 1:16, 2:8}[l], \ |
| 114 | num_anchors//3, num_classes+5)) for l in range(3)] |
| 115 | |
| 116 | model_body = yolo_body(image_input, num_anchors//3, num_classes) |
| 117 | print('Create YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes)) |
| 118 | |
| 119 | if load_pretrained: |
| 120 | model_body.load_weights(weights_path, by_name=True, skip_mismatch=True) |
| 121 | print('Load weights {}.'.format(weights_path)) |
| 122 | if freeze_body in [1, 2]: |
| 123 | # Freeze darknet53 body or freeze all but 3 output layers. |
| 124 | num = (185, len(model_body.layers)-3)[freeze_body-1] |
| 125 | for i in range(num): model_body.layers[i].trainable = False |
| 126 | print('Freeze the first {} layers of total {} layers.'.format(num, len(model_body.layers))) |
| 127 | |
| 128 | model_loss = Lambda(yolo_loss, output_shape=(1,), name='yolo_loss', |
| 129 | arguments={'anchors': anchors, 'num_classes': num_classes, 'ignore_thresh': 0.5})( |
| 130 | [*model_body.output, *y_true]) |
| 131 | model = Model([model_body.input, *y_true], model_loss) |
| 132 | |
| 133 | return model |
| 134 | |
| 135 | def create_tiny_model(input_shape, anchors, num_classes, load_pretrained=True, freeze_body=2, |
| 136 | weights_path='model_data/tiny_yolo_weights.h5'): |