Source code for examples.model_zoo.yolo

#! /usr/bin/python
# -*- coding: utf-8 -*-
"""YOLOv4 for MS-COCO.

# Reference:
- [tensorflow-yolov4-tflite](
    https://github.com/hunglc007/tensorflow-yolov4-tflite)

"""

import numpy as np
import tensorlayerx as tlx
from tensorlayerx.nn import Mish
from tensorlayerx.nn import Conv2d, MaxPool2d, BatchNorm2d, ZeroPad2d, UpSampling2d, Concat, Elementwise
from tensorlayerx.nn import Module, Sequential
from tensorlayerx import logging

__all__ = ['YOLOv4']

INPUT_SIZE = 416
weights_url = {'link': 'https://pan.baidu.com/s/1MC1dmEwpxsdgHO1MZ8fYRQ', 'password': 'idsz'}


class Convolutional(Module):
    """
    Create Convolution layer
    Because it is only a stack of reference layers, there is no build, so self._built=True
    """

    def __init__(self, filters_shape, downsample=False, activate=True, bn=True, activate_type='leaky', name=None):
        super(Convolutional, self).__init__()
        self.act = activate
        self.act_type = activate_type
        self.downsample = downsample
        self.bn = bn
        self._built = True
        if downsample:
            padding = 'VALID'
            strides = 2
        else:
            strides = 1
            padding = 'SAME'

        if bn:
            b_init = None
        else:
            b_init = tlx.nn.initializers.constant(value=0.0)

        self.zeropad = ZeroPad2d(((1, 0), (1, 0)))
        self.conv = Conv2d(
            out_channels=filters_shape[-1], in_channels=filters_shape[2], kernel_size=(filters_shape[0], filters_shape[1]),
            stride=(strides, strides), padding=padding, b_init=b_init, name=name
        )

        if bn:
            if activate ==True:
                if activate_type == 'leaky':
                    self.batchnorm2d = BatchNorm2d(act='leaky_relu0.1', num_features=filters_shape[-1])
                elif activate_type == 'mish':
                    self.batchnorm2d = BatchNorm2d(act=Mish, num_features=filters_shape[-1])
            else:
                self.batchnorm2d = BatchNorm2d(act=None, num_features=filters_shape[-1])

    def forward(self, input):
        if self.downsample:
            input = self.zeropad(input)

        output = self.conv(input)

        if self.bn:
            output = self.batchnorm2d(output)
        return output


class residual_block(Module):

    def __init__(self, input_channel, filter_num1, filter_num2, activate_type='leaky'):
        super(residual_block, self).__init__()
        self.conv1 = Convolutional(filters_shape=(1, 1, input_channel, filter_num1), activate_type=activate_type)
        self.conv2 = Convolutional(filters_shape=(3, 3, filter_num1, filter_num2), activate_type=activate_type)
        self.add = Elementwise(tlx.add)

    def forward(self, inputs):
        output = self.conv1(inputs)
        output = self.conv2(output)
        output = self.add([inputs, output])
        return output


def residual_block_num(num, input_channel, filter_num1, filter_num2, activate_type='leaky'):
    residual_list = []
    for i in range(num):
        residual_list.append(residual_block(input_channel, filter_num1, filter_num2, activate_type=activate_type))
    return Sequential(residual_list)


class cspdarknet53(Module):

    def __init__(self):
        super(cspdarknet53, self).__init__()
        self._built = True
        self.conv1_1 = Convolutional((3, 3, 3, 32), activate_type='mish')
        self.conv1_2 = Convolutional((3, 3, 32, 64), downsample=True, activate_type='mish')
        self.conv1_3 = Convolutional((1, 1, 64, 64), activate_type='mish', name='conv_rote_block_1')
        self.conv1_4 = Convolutional((1, 1, 64, 64), activate_type='mish')
        self.residual_1 = residual_block_num(1, 64, 32, 64, activate_type="mish")

        self.conv2_1 = Convolutional((1, 1, 64, 64), activate_type='mish')
        self.concat = Concat()
        self.conv2_2 = Convolutional((1, 1, 128, 64), activate_type='mish')
        self.conv2_3 = Convolutional((3, 3, 64, 128), downsample=True, activate_type='mish')
        self.conv2_4 = Convolutional((1, 1, 128, 64), activate_type='mish', name='conv_rote_block_2')
        self.conv2_5 = Convolutional((1, 1, 128, 64), activate_type='mish')
        self.residual_2 = residual_block_num(2, 64, 64, 64, activate_type='mish')

        self.conv3_1 = Convolutional((1, 1, 64, 64), activate_type='mish')
        self.conv3_2 = Convolutional((1, 1, 128, 128), activate_type='mish')
        self.conv3_3 = Convolutional((3, 3, 128, 256), downsample=True, activate_type='mish')
        self.conv3_4 = Convolutional((1, 1, 256, 128), activate_type='mish', name='conv_rote_block_3')
        self.conv3_5 = Convolutional((1, 1, 256, 128), activate_type='mish')
        self.residual_3 = residual_block_num(8, 128, 128, 128, activate_type="mish")

        self.conv4_1 = Convolutional((1, 1, 128, 128), activate_type='mish')
        self.conv4_2 = Convolutional((1, 1, 256, 256), activate_type='mish')
        self.conv4_3 = Convolutional((3, 3, 256, 512), downsample=True, activate_type='mish')
        self.conv4_4 = Convolutional((1, 1, 512, 256), activate_type='mish', name='conv_rote_block_4')
        self.conv4_5 = Convolutional((1, 1, 512, 256), activate_type='mish')
        self.residual_4 = residual_block_num(8, 256, 256, 256, activate_type="mish")

        self.conv5_1 = Convolutional((1, 1, 256, 256), activate_type='mish')
        self.conv5_2 = Convolutional((1, 1, 512, 512), activate_type='mish')
        self.conv5_3 = Convolutional((3, 3, 512, 1024), downsample=True, activate_type='mish')
        self.conv5_4 = Convolutional((1, 1, 1024, 512), activate_type='mish', name='conv_rote_block_5')
        self.conv5_5 = Convolutional((1, 1, 1024, 512), activate_type='mish')
        self.residual_5 = residual_block_num(4, 512, 512, 512, activate_type="mish")

        self.conv6_1 = Convolutional((1, 1, 512, 512), activate_type='mish')
        self.conv6_2 = Convolutional((1, 1, 1024, 1024), activate_type='mish')
        self.conv6_3 = Convolutional((1, 1, 1024, 512))
        self.conv6_4 = Convolutional((3, 3, 512, 1024))
        self.conv6_5 = Convolutional((1, 1, 1024, 512))

        self.maxpool1 = MaxPool2d(kernel_size=(13, 13), stride=(1, 1))
        self.maxpool2 = MaxPool2d(kernel_size=(9, 9), stride=(1, 1))
        self.maxpool3 = MaxPool2d(kernel_size=(5, 5), stride=(1, 1))

        self.conv7_1 = Convolutional((1, 1, 2048, 512))
        self.conv7_2 = Convolutional((3, 3, 512, 1024))
        self.conv7_3 = Convolutional((1, 1, 1024, 512))

    def forward(self, input_data):
        input_data = self.conv1_1(input_data)
        input_data = self.conv1_2(input_data)
        route = input_data
        route = self.conv1_3(route)
        input_data = self.conv1_4(input_data)
        input_data = self.residual_1(input_data)

        input_data = self.conv2_1(input_data)
        input_data = self.concat([input_data, route])
        input_data = self.conv2_2(input_data)
        input_data = self.conv2_3(input_data)
        route = input_data
        route = self.conv2_4(route)
        input_data = self.conv2_5(input_data)
        input_data = self.residual_2(input_data)

        input_data = self.conv3_1(input_data)
        input_data = self.concat([input_data, route])
        input_data = self.conv3_2(input_data)
        input_data = self.conv3_3(input_data)
        route = input_data
        route = self.conv3_4(route)
        input_data = self.conv3_5(input_data)
        input_data = self.residual_3(input_data)

        input_data = self.conv4_1(input_data)
        input_data = self.concat([input_data, route])
        input_data = self.conv4_2(input_data)
        route_1 = input_data
        input_data = self.conv4_3(input_data)
        route = input_data
        route = self.conv4_4(route)
        input_data = self.conv4_5(input_data)
        input_data = self.residual_4(input_data)

        input_data = self.conv5_1(input_data)
        input_data = self.concat([input_data, route])
        input_data = self.conv5_2(input_data)
        route_2 = input_data
        input_data = self.conv5_3(input_data)
        route = input_data
        route = self.conv5_4(route)
        input_data = self.conv5_5(input_data)
        input_data = self.residual_5(input_data)

        input_data = self.conv6_1(input_data)
        input_data = self.concat([input_data, route])

        input_data = self.conv6_2(input_data)
        input_data = self.conv6_3(input_data)
        input_data = self.conv6_4(input_data)
        input_data = self.conv6_5(input_data)

        maxpool1 = self.maxpool1(input_data)
        maxpool2 = self.maxpool2(input_data)
        maxpool3 = self.maxpool3(input_data)
        input_data = self.concat([maxpool1, maxpool2, maxpool3, input_data])

        input_data = self.conv7_1(input_data)
        input_data = self.conv7_2(input_data)
        input_data = self.conv7_3(input_data)

        return route_1, route_2, input_data


class YOLOv4_model(Module):

    def __init__(self, NUM_CLASS):
        super(YOLOv4_model, self).__init__()
        self.cspdarnnet = cspdarknet53()

        self.conv1_1 = Convolutional((1, 1, 512, 256))
        self.upsamle = UpSampling2d(scale=2)
        self.conv1_2 = Convolutional((1, 1, 512, 256), name='conv_yolo_1')
        self.concat = Concat()

        self.conv2_1 = Convolutional((1, 1, 512, 256))
        self.conv2_2 = Convolutional((3, 3, 256, 512))
        self.conv2_3 = Convolutional((1, 1, 512, 256))
        self.conv2_4 = Convolutional((3, 3, 256, 512))
        self.conv2_5 = Convolutional((1, 1, 512, 256))

        self.conv3_1 = Convolutional((1, 1, 256, 128))
        self.conv3_2 = Convolutional((1, 1, 256, 128), name='conv_yolo_2')

        self.conv4_1 = Convolutional((1, 1, 256, 128))
        self.conv4_2 = Convolutional((3, 3, 128, 256))
        self.conv4_3 = Convolutional((1, 1, 256, 128))
        self.conv4_4 = Convolutional((3, 3, 128, 256))
        self.conv4_5 = Convolutional((1, 1, 256, 128))

        self.conv5_1 = Convolutional((3, 3, 128, 256), name='conv_route_1')
        self.conv5_2 = Convolutional((1, 1, 256, 3 * (NUM_CLASS + 5)), activate=False, bn=False)

        self.conv6_1 = Convolutional((3, 3, 128, 256), downsample=True, name='conv_route_2')
        self.conv6_2 = Convolutional((1, 1, 512, 256))
        self.conv6_3 = Convolutional((3, 3, 256, 512))
        self.conv6_4 = Convolutional((1, 1, 512, 256))
        self.conv6_5 = Convolutional((3, 3, 256, 512))
        self.conv6_6 = Convolutional((1, 1, 512, 256))

        self.conv7_1 = Convolutional((3, 3, 256, 512), name='conv_route_3')
        self.conv7_2 = Convolutional((1, 1, 512, 3 * (NUM_CLASS + 5)), activate=False, bn=False)
        self.conv7_3 = Convolutional((3, 3, 256, 512), downsample=True, name='conv_route_4')

        self.conv8_1 = Convolutional((1, 1, 1024, 512))
        self.conv8_2 = Convolutional((3, 3, 512, 1024))
        self.conv8_3 = Convolutional((1, 1, 1024, 512))
        self.conv8_4 = Convolutional((3, 3, 512, 1024))
        self.conv8_5 = Convolutional((1, 1, 1024, 512))

        self.conv9_1 = Convolutional((3, 3, 512, 1024))
        self.conv9_2 = Convolutional((1, 1, 1024, 3 * (NUM_CLASS + 5)), activate=False, bn=False)

    def forward(self, inputs):
        route_1, route_2, conv = self.cspdarnnet(inputs)

        route = conv
        conv = self.conv1_1(conv)
        conv = self.upsamle(conv)
        route_2 = self.conv1_2(route_2)
        conv = self.concat([route_2, conv])

        conv = self.conv2_1(conv)
        conv = self.conv2_2(conv)
        conv = self.conv2_3(conv)
        conv = self.conv2_4(conv)
        conv = self.conv2_5(conv)

        route_2 = conv
        conv = self.conv3_1(conv)
        conv = self.upsamle(conv)
        route_1 = self.conv3_2(route_1)
        conv = self.concat([route_1, conv])

        conv = self.conv4_1(conv)
        conv = self.conv4_2(conv)
        conv = self.conv4_3(conv)
        conv = self.conv4_4(conv)
        conv = self.conv4_5(conv)

        route_1 = conv
        conv = self.conv5_1(conv)
        conv_sbbox = self.conv5_2(conv)

        conv = self.conv6_1(route_1)
        conv = self.concat([conv, route_2])

        conv = self.conv6_2(conv)
        conv = self.conv6_3(conv)
        conv = self.conv6_4(conv)
        conv = self.conv6_5(conv)
        conv = self.conv6_6(conv)

        route_2 = conv
        conv = self.conv7_1(conv)
        conv_mbbox = self.conv7_2(conv)
        conv = self.conv7_3(route_2)
        conv = self.concat([conv, route])

        conv = self.conv8_1(conv)
        conv = self.conv8_2(conv)
        conv = self.conv8_3(conv)
        conv = self.conv8_4(conv)
        conv = self.conv8_5(conv)

        conv = self.conv9_1(conv)
        conv_lbbox = self.conv9_2(conv)

        return conv_sbbox, conv_mbbox, conv_lbbox


[docs]def YOLOv4(NUM_CLASS, pretrained=False): """Pre-trained YOLOv4 model. Parameters ------------ NUM_CLASS : int Number of classes in final prediction. pretrained : boolean Whether to load pretrained weights. Default False. Examples --------- Object Detection with YOLOv4, see `computer_vision.py <https://github.com/tensorlayer/tensorlayer/blob/master/tensorlayer/app/computer_vision.py>`__ With TensorLayer >>> # get the whole model, without pre-trained YOLOv4 parameters >>> yolov4 = YOLOv4(NUM_CLASS=80, pretrained=False) >>> # get the whole model, restore pre-trained YOLOv4 parameters >>> yolov4 = YOLOv4(NUM_CLASS=80, pretrained=True) >>> # use for inferencing >>> output = yolov4(img) """ network = YOLOv4_model(NUM_CLASS=NUM_CLASS) if pretrained: restore_params(network, model_path='model/yolov4_model.npz') return network
def restore_params(network, model_path='model.npz'): logging.info("Restore pre-trained weights") try: npz = np.load(model_path, allow_pickle=True) except: print("Download the model file, placed in the /model ") print("Weights download: ", weights_url['link'], "password:", weights_url['password']) txt_path = 'model/yolov4_weights3_config.txt' f = open(txt_path, "r") line = f.readlines() for i in range(len(line)): network.all_weights[i].assign(npz[line[i].strip()]) logging.info(" Loading weights %s in %s" % (network.all_weights[i].shape, network.all_weights[i].name)) def tl2_weights_to_tl3_weights( weights_2_path='model/weights_2.txt', weights_3_path='model/weights_3.txt', txt_path='model/yolov4_weights_config.txt' ): weights_2_path = weights_2_path weights_3_path = weights_3_path txt_path = txt_path f1 = open(weights_2_path, "r") f2 = open(weights_3_path, "r") f3 = open(txt_path, "r") line1 = f1.readlines() line2 = f2.readlines() line3 = f3.readlines() _dicts = {} for i in range(len(line1)): _dicts[line1[i].strip()] = line3[i].strip() for j in range(len(line2)): print(_dicts[line2[j].strip()])