load_weights copy.py

"""
This file contains the code to load parsed weights that are in the DarkNet
format into TensorFlow layers
"""
import itertools
from tensorflow import keras as ks
from yolo.modeling.building_blocks import DarkConv


def split_converter(lst, i, j=None):
    if j is None:
        return lst.data[:i], lst.data[i:]
    return lst.data[:i], lst.data[i:j], lst.data[j:]


def interleve_weights(block):
    """merge weights to fit the DarkResnet block style"""
    if len(block) == 0:
        return []
    weights_temp = []
    for layer in block:
        weights = layer.get_weights()
        weights = [tuple(weights[0:3]), tuple(weights[3:])]
        weights_temp.append(weights)
    top, bottom = tuple(zip(*weights_temp))
    weights = list(itertools.chain.from_iterable(top)) + \
        list(itertools.chain.from_iterable(bottom))
    return weights


def get_darknet53_tf_format(net, only_weights=True):
    """convert weights from darknet sequntial to tensorflow weave, Darknet53 Backbone"""
    combo_blocks = []
    for i in range(2):
        layer = net.pop(0)
        combo_blocks.append(layer)
    # ugly code i will document, very tired
    encoder = []
    while len(net) != 0:
        blocks = []
        layer = net.pop(0)
        while layer._type != "shortcut":
            blocks.append(layer)
            layer = net.pop(0)
        encoder.append(blocks)
    new_net = combo_blocks + encoder
    weights = []
    if only_weights:
        for block in new_net:
            if type(block) != list:
                weights.append(block.get_weights())
            else:
                weights.append(interleve_weights(block))
    print("converted/interleved weights for tensorflow format")
    return new_net, weights


def get_tiny_tf_format(encoder):
    weights = []
    for layer in encoder:
        if layer._type != "maxpool":
            weights.append(layer.get_weights())
    return encoder, weights


def load_weights_dnBackbone(backbone, encoder, mtype="darknet53"):
    # get weights for backbone
    if mtype == "darknet53":
        encoder, weights_encoder = get_darknet53_tf_format(encoder[:])
    elif mtype == "darknet_tiny":
        encoder, weights_encoder = get_tiny_tf_format(encoder[:])

    # set backbone weights
    print(
        f"\nno. layers: {len(backbone.layers)}, no. weights: {len(weights_encoder)}"
    )
    set_darknet_weights(backbone, weights_encoder)

    backbone.trainable = False
    print(f"\nsetting backbone.trainable to: {backbone.trainable}\n")
    return


def load_weights_dnHead(head, decoder, v4=True):
    # get weights for head
    decoder, weights_decoder, head_layers, head_weights = get_decoder_weights(
        decoder)
    # set detection head weights
    print(
        f"\nno. layers: {len(head.layers)}, no. weights: {len(weights_decoder)}"
    )
    flat_full = list(flatten_model(head, r_list=False))
    flat_main = flat_full[:-3]
    flat_head = flat_full[-3:]

    # not the right way to do it
    if v4:
        flat_main.insert(1, flat_main[-1])

    print(len(flat_main), len(decoder))
    print(len(flat_head), len(head_layers))

    set_darknet_weights(head, weights_decoder, flat_model=flat_main)
    set_darknet_weights_head(flat_head, head_weights)

    head.trainable = False
    print(f"\nsetting head.trainable to: {head.trainable}\n")
    return


# DEBUGGING
def print_layer_shape(layer):
    try:
        weights = layer.get_weights()
    except:
        weights = layer
    for item in weights:
        print(item.shape)
    return


def flatten_model(model, r_list=True):
    for layer in model.layers:
        if r_list and isinstance(model, ks.Model):
            yield from model.layers
        else:
            yield layer


def set_darknet_weights_head(flat_head, weights_head):
    for layer in flat_head:
        weights = layer.get_weights()
        for weight in weights:
            print(weight.shape)
        weight_depth = weights[0].shape[-2]
        for weight in weights_head:
            if weight[0].shape[-2] == weight_depth:
                print(
                    f"loaded weights for layer: head layer with depth {weight_depth}  -> name: {layer.name}",
                    sep='      ',
                    end="\r")
                layer.set_weights(weight)
    return


def set_darknet_weights(model, weights_list, flat_model=None):
    if flat_model == None:
        zip_fill = flatten_model(model)
    else:
        zip_fill = flat_model
    for i, (layer, weights) in enumerate(zip(zip_fill, weights_list)):
        print(layer.name, len(weights))
        #layer.set_weights(weights)
    return


def split_decoder(lst):
    decoder = []
    outputs = []
    for layer in lst:
        if layer._type == 'yolo':
            outputs.append(decoder.pop())
            outputs.append(layer)
        else:
            decoder.append(layer)
    return decoder, outputs


def get_decoder_weights(decoder):
    layers = [[]]
    block = []
    weights = []

    decoder, head = split_decoder(decoder)

    # get decoder weights and group them together
    for i, layer in enumerate(decoder):
        if layer._type == "route" and len(
                layer.layers) >= 2 and decoder[i - 1]._type != 'maxpool':
            layers.append([])
            layers.append(block)
            block = []
        elif layer._type == "route" and decoder[i - 1]._type != 'maxpool':
            layers.append(block)
            block = []
        elif (layer._type == "route" and decoder[i - 1]._type
              == "maxpool") or layer._type == "maxpool":
            # made only for spp
            continue
        elif layer._type == "convolutional":
            block.append(layer)
        # else:
        #     # if you upsample
        #     layers.append([])

    if len(block) > 0:
        layers.append(block)

    # interleve weights for blocked layers
    for layer in layers:
        weights.append(interleve_weights(layer))

    # get weights for output detection heads
    head_weights = []
    head_layers = []
    for layer in (head):
        if layer != None and layer._type == "convolutional":
            head_weights.append(layer.get_weights())
            head_layers.append(layer)

    return layers, weights, head_layers, head_weights