Esempio n. 1
0
def darknet_base(inputs, include_yolo_head=True):
    """
    Builds Darknet53 by reading the YOLO configuration file

    :param inputs: Input tensor
    :param include_yolo_head: Includes the YOLO head
    :return: A list of output layers and the network config
    """
    path = os.path.join(ROOT_DIR, 'cfg', '{}.cfg'.format(YOLO_VERSION))
    blocks = Parser.parse_cfg(path)
    x, layers, yolo_layers = inputs, [], []
    ptr = 0
    config = {}

    for block in blocks:
        block_type = block['type']

        if block_type == 'net':
            config = _read_net_config(block)

        elif block_type == 'convolutional':
            x, layers, yolo_layers, ptr = _build_conv_layer(
                x, block, layers, yolo_layers, ptr, config)

        elif block_type == 'shortcut':
            x, layers, yolo_layers, ptr = _build_shortcut_layer(
                x, block, layers, yolo_layers, ptr)

        elif block_type == 'yolo':
            x, layers, yolo_layers, ptr = _build_yolo_layer(
                x, block, layers, yolo_layers, ptr, config)

        elif block_type == 'route':
            x, layers, yolo_layers, ptr = _build_route_layer(
                x, block, layers, yolo_layers, ptr)

        elif block_type == 'upsample':
            x, layers, yolo_layers, ptr = _build_upsample_layer(
                x, block, layers, yolo_layers, ptr)

        elif block_type == 'maxpool':
            x, layers, yolo_layers, ptr = _build_maxpool_layer(
                x, block, layers, yolo_layers, ptr)

        else:
            raise ValueError(
                '{} not recognized as block type'.format(block_type))

    _verify_weights_completed_consumed(ptr)

    if include_yolo_head:
        output_layers = yolo_layers
        return tf.keras.layers.Concatenate(axis=1)(output_layers), config
    else:
        output_layers = [
            layers[i - 1] for i in range(len(layers)) if layers[i] is None
        ]

        # NOTE: Apparently TFLite doesn't like Concatenate.
        # return tf.keras.layers.Concatenate(axis=1)(output_layers), config
        return output_layers, config
Esempio n. 2
0
def darknet_base(inputs,
                 yolo_version='yolov3',
                 data_dir='./weights',
                 include_yolo_head=True):
    """
    Builds Darknet53 by reading the YOLO configuration file

    :param inputs: Input tensor
    :param include_yolo_head: Includes the YOLO head
    :return: A list of output layers and the network config
    """
    weights_file = open(
        os.path.join(data_dir, '{}.weights'.format(yolo_version)), 'rb')
    major, minor, revision = np.ndarray(shape=(3, ),
                                        dtype='int32',
                                        buffer=weights_file.read(12))
    if (major * 10 + minor) >= 2 and major < 1000 and minor < 1000:
        seen = np.ndarray(shape=(1, ),
                          dtype='int64',
                          buffer=weights_file.read(8))
    else:
        seen = np.ndarray(shape=(1, ),
                          dtype='int32',
                          buffer=weights_file.read(4))
    print('Weights Header: ', major, minor, revision, seen)

    path = os.path.join(data_dir, '{}.cfg'.format(yolo_version))
    blocks = Parser.parse_cfg(path)
    x, layers, yolo_layers = inputs, [], []
    ptr = 0
    config = {}

    for block in blocks:
        block_type = block['type']

        if block_type == 'net':
            config = _read_net_config(block, data_dir, yolo_version)

        elif block_type == 'convolutional':
            x, layers, yolo_layers, ptr = _build_conv_layer(
                x, block, layers, yolo_layers, ptr, config, weights_file)

        elif block_type == 'shortcut':
            x, layers, yolo_layers, ptr = _build_shortcut_layer(
                x, block, layers, yolo_layers, ptr)

        elif block_type == 'yolo':
            x, layers, yolo_layers, ptr = _build_yolo_layer(
                x, block, layers, yolo_layers, ptr, config)

        elif block_type == 'route':
            x, layers, yolo_layers, ptr = _build_route_layer(
                x, block, layers, yolo_layers, ptr)

        elif block_type == 'upsample':
            x, layers, yolo_layers, ptr = _build_upsample_layer(
                x, block, layers, yolo_layers, ptr)

        elif block_type == 'maxpool':
            x, layers, yolo_layers, ptr = _build_maxpool_layer(
                x, block, layers, yolo_layers, ptr)

        else:
            raise ValueError(
                '{} not recognized as block type'.format(block_type))

    _verify_weights_completed_consumed(ptr, weights_file)

    if include_yolo_head:
        output_layers = yolo_layers
        return tf.keras.layers.Concatenate(axis=1)(output_layers), config
    else:
        output_layers = [
            layers[i - 1] for i in range(len(layers)) if layers[i] is None
        ]

        # NOTE: Apparently TFLite doesn't like Concatenate.
        # return tf.keras.layers.Concatenate(axis=1)(output_layers), config
        return output_layers, config