LP_detection.py

#! usr/bin/python
import argparse
import datetime
import math
import numpy as np
import os
import sys
import yaml

import mxboard
import mxnet
from mxnet import nd
from mxnet.gluon import nn
from gluoncv.model_zoo.densenet import _make_dense_block, _make_transition

from yolo_modules import global_variable
from yolo_modules import licence_plate_render
from yolo_modules import yolo_cv
from yolo_modules import yolo_gluon

os.environ['MXNET_ENABLE_GPU_P2P'] = '0'
os.environ['MXNET_CUDNN_AUTOTUNE_DEFAULT'] = '0'


def main():
    args = Parser()
    LP_detection = LicencePlateDetectioin(args)
    available_mode = ['train', 'valid', 'export']
    assert args.mode in available_mode, \
        'Available Modes Are {}'.format(available_mode)

    exec "LP_detection.%s()" % args.mode


def Parser():
    parser = argparse.ArgumentParser(prog="python LP_detection.py")
    parser.add_argument("version", help="v2 is good")
    parser.add_argument("mode", help="train/valid/video")

    parser.add_argument("--gpu", help="gpu index", dest="gpu", default="0")

    # -------------------- Train/Valid -------------------- #
    parser.add_argument(
        "--weight",
        dest="weight", default=None,
        help="pretrain weight file")

    parser.add_argument(
        "--record",
        dest="record", default=1, type=int,
        help="record to tensorboard or not while training")

    # ----------------------- Video ----------------------- #
    parser = yolo_cv.add_video_parser(parser)

    return parser.parse_args()


class LPDenseNet(mxnet.gluon.HybridBlock):
    # copy from:
    # https://github.com/dmlc/gluon-cv/blob/3658339acbdfc78c2191c687e4430e3a673
    # 66b7d/gluoncv/model_zoo/densenet.py#L620
    # Densely Connected Convolutional Networks
    # <https://arxiv.org/pdf/1608.06993.pdf>
    def __init__(self, num_init_features, growth_rate, block_config,
                 bn_size=4, dropout=0, classes=1, **kwargs):
        super(LPDenseNet, self).__init__(**kwargs)
        with self.name_scope():
            self.features = nn.HybridSequential(prefix='')
            self.features.add(nn.Conv2D(num_init_features, kernel_size=7,
                                        strides=2, padding=3, use_bias=False))
            self.features.add(nn.BatchNorm())
            self.features.add(nn.Activation('relu'))
            self.features.add(nn.MaxPool2D(pool_size=3, strides=2, padding=1))
            # Add dense blocks
            num_features = num_init_features
            for i, num_layers in enumerate(block_config):
                self.features.add(
                    _make_dense_block(
                        num_layers, bn_size, growth_rate, dropout, i+1))

                num_features = num_features + num_layers * growth_rate
                if i != len(block_config) - 1:
                    self.features.add(_make_transition(num_features // 2))
                    num_features = num_features // 2
            self.features.add(nn.BatchNorm())
            self.features.add(nn.Activation('relu'))

            self.features.add(nn.Conv2D(512, (3, 3), padding=1))
            self.features.add(nn.BatchNorm())
            self.features.add(nn.Activation('relu'))

            self.features.add(nn.Conv2D(7+classes, (1, 1)))

    def hybrid_forward(self, F, x):
        x = self.features(x)
        return x


class LicencePlateDetectioin():
    def __init__(self, args):
        spec_path = os.path.join(args.version, 'spec.yaml')
        with open(spec_path) as f:
            spec = yaml.load(f)

        for key in spec:
            setattr(self, key, spec[key])

        self.version = args.version
        self.ctx = yolo_gluon.get_ctx(args.gpu)
        self.export_file = args.version + '/export/'
        self.num_downsample = len(self.block_config) + 1

        assert args.mode in ['train', 'valid', 'export', 'video']
        if args.mode == 'video':
            return

        self.backup_dir = os.path.join(args.version, 'backup')
        self.net = LPDenseNet(
            self.num_init_features,
            self.growth_rate,
            self.block_config,
            classes=self.LP_num_class)

        if args.weight is None:
            args.weight = yolo_gluon.get_latest_weight_from(
                self.backup_dir)

        yolo_gluon.init_NN(self.net, args.weight, self.ctx)

        if args.mode == 'train':
            self.record = args.record
            self._init_train()

    def train(self):

        self._train_or_valid('train')

    def valid(self):

        self._train_or_valid('val')

    def export(self):
        shape = (1, 3, self.size[0], self.size[1])
        yolo_gluon.export(self.net, shape, self.ctx[0], self.export_file, onnx=1, epoch=0)

    def predict_LP(self, batch_out):
        use_np = True if type(batch_out) == np.ndarray else False
        out = batch_out.transpose((0, 2, 3, 1))[0]
        # (10L, 16L, features)
        best_index = out[:, :, 0].reshape(-1).argmax(axis=0)
        out = out.reshape((-1, self.LP_slice_point[-1]))

        pred = out[best_index].asnumpy()[0] if not use_np else out[best_index]

        pred[0] = yolo_gluon.np_sigmoid(pred[0])
        pred[1:4] *= 1000
        for i in range(3):
            p = (yolo_gluon.np_sigmoid(pred[i+4]) - 0.5) * 2 * self.LP_r_max[i]
            pred[i+4] = p * math.pi / 180.

        return pred

    def slice_out(self, x, use_np=False):
        x = x.transpose((0, 2, 3, 1))
        i = 0
        outs = []
        for pt in self.LP_slice_point:
            if use_np:
                y = x[:, :, :, i:pt]
            else:
                y = x.slice_axis(begin=i, end=pt, axis=-1)

            outs.append(y)
            i = pt

        return outs

    # -------------------- Train -------------------- #
    def _train_or_valid(self, mode):
        print(global_variable.cyan)
        print(mode)
        print(global_variable.reset_color)
        if mode == 'val':
            self.batch_size = 1
            ax = yolo_cv.init_matplotlib_figure()
            # self.net = yolo_gluon.init_executor(
            #    self.export_file, self.size, self.ctx[0])

        # -------------------- background -------------------- #
        LP_generator = licence_plate_render.LPGenerator(*self.size)
        bg_iter = yolo_gluon.load_background(mode, self.batch_size, *self.size)

        # -------------------- main loop -------------------- #
        self.backward_counter = 0
        while True:
            if (self.backward_counter % 3 == 0 or 'bg' not in locals()):
                bg = yolo_gluon.ImageIter_next_batch(bg_iter)
                bg = bg.as_in_context(self.ctx[0]) / 255.

            # -------------------- render dataset -------------------- #
            imgs, labels = LP_generator.add(bg, self.LP_r_max, add_rate=0.5)

            if mode == 'train':
                batch_xs = yolo_gluon.split_render_data(imgs, self.ctx)
                batch_ys = yolo_gluon.split_render_data(labels, self.ctx)
                self._train_batch_LP(batch_xs, batch_ys)

            elif mode == 'val':
                batch_out = self.net(imgs)
                pred = self.predict_LP(batch_out)

                img = yolo_gluon.batch_ndimg_2_cv2img(imgs)[0]
                labels = labels.asnumpy()

                img, _ = LP_generator.project_rect_6d.add_edges(
                    img, labels[0, 0, 1:7])

                img, clipped_LP = LP_generator.project_rect_6d.add_edges(
                    img, pred[1:])

                yolo_cv.matplotlib_show_img(ax, img)
                print('Labels: {}'.format(labels))
                print('predictions: {}'.format(pred))
                raw_input('--------------------------------------------------')

    def _init_train(self):
        self.exp = self.exp + datetime.datetime.now().strftime("%m-%dx%H-%M")

        self.batch_size *= len(self.ctx)

        print(global_variable.yellow)
        print('Batch Size = {}'.format(self.batch_size))
        print('Record Step = {}'.format(self.record_step))

        #self.L1_loss = mxnet.gluon.loss.L1Loss()
        #self.L2_loss = mxnet.gluon.loss.L2Loss()
        self.HB_loss = mxnet.gluon.loss.HuberLoss()
        self.LG_loss = mxnet.gluon.loss.LogisticLoss(label_format='binary')
        self.CE_loss = mxnet.gluon.loss.SoftmaxCrossEntropyLoss(
            from_logits=False, sparse_label=False)

        # -------------------- init trainer -------------------- #
        optimizer = mxnet.optimizer.create(
            'adam',
            learning_rate=self.learning_rate,
            multi_precision=False)

        self.trainer = mxnet.gluon.Trainer(
            self.net.collect_params(),
            optimizer=optimizer)

        logdir = os.path.join(self.version, 'logs')
        self.sw = mxboard.SummaryWriter(logdir=logdir, verbose=True)

        if not os.path.exists(self.backup_dir):
            os.makedirs(self.backup_dir)

    def _find_best_LP(self, L, gpu_index):
        step = 2**self.num_downsample
        h_feature_max = self.size[0] / 2**self.num_downsample - 1
        w_feature_max = self.size[1] / 2**self.num_downsample - 1
        h_feature = np.clip(int(L[8].asnumpy()/step), 0, h_feature_max)
        w_feature = np.clip(int(L[7].asnumpy()/step), 0, w_feature_max)
        #print(h_feature_max, w_feature_max)
        #print(h_feature, w_feature)
        t_X = L[1] / 1000.
        t_Y = L[2] / 1000.
        t_Z = L[3] / 1000.

        r1_max = self.LP_r_max[0] * math.pi / 180.
        r2_max = self.LP_r_max[1] * math.pi / 180.
        r3_max = self.LP_r_max[2] * math.pi / 180.  # 2 * theta (rad)

        t_r1 = yolo_gluon.nd_inv_sigmoid(L[4] / r1_max / 2. + 0.5)
        t_r2 = yolo_gluon.nd_inv_sigmoid(L[5] / r2_max / 2. + 0.5)
        t_r3 = yolo_gluon.nd_inv_sigmoid(L[6] / r3_max / 2. + 0.5)

        label = nd.concat(t_X, t_Y, t_Z, t_r1, t_r2, t_r3, dim=-1)

        return (h_feature, w_feature), label

    def _loss_mask_LP(self, label_batch, gpu_index):
        """Generate training targets given predictions and label_batch.
        label_batch: bs*object*[class, cent_y, cent_x, box_h, box_w, rotate]
        """
        bs = label_batch.shape[0]
        ctx = self.ctx[gpu_index]
        h_ = self.size[0] / 2**self.num_downsample
        w_ = self.size[1] / 2**self.num_downsample

        score = nd.zeros((bs, h_, w_, 1), ctx=ctx)
        mask = nd.zeros((bs, h_, w_, 1), ctx=ctx)
        pose_xy = nd.zeros((bs, h_, w_, 2), ctx=ctx)
        pose_z = nd.zeros((bs, h_, w_, 1), ctx=ctx)
        pose_r = nd.zeros((bs, h_, w_, 3), ctx=ctx)
        LP_class = nd.zeros((bs, h_, w_, self.LP_num_class), ctx=ctx)

        for b in range(bs):
            for L in label_batch[b]:  # all object in the image
                if L[0] < 0:
                    continue

                else:
                    (h_f, w_f), p_6D = self._find_best_LP(L, gpu_index)
                    score[b, h_f, w_f, :] = 1.0  # others are zero
                    mask[b, h_f, w_f, :] = 1.0  # others are zero
                    pose_xy[b, h_f, w_f, :] = p_6D[:2]
                    pose_z[b, h_f, w_f, :] = p_6D[2]
                    pose_r[b, h_f, w_f, :] = p_6D[3:]
                    LP_class[b, h_f, w_f, L[-1]] = 1

        return [score, pose_xy, pose_z, pose_r, LP_class], mask

    def _train_batch_LP(self, bxs, bys):
        all_gpu_loss = [None] * len(self.ctx)
        with mxnet.autograd.record():
            for gpu_i, (bx, by) in enumerate(zip(bxs, bys)):
                # ---------- Forward ---------- #
                by_ = self.net(bx)
                by_ = self.slice_out(by_)

                with mxnet.autograd.pause():
                    by, mask = self._loss_mask_LP(by, gpu_i)
                    ones = nd.ones_like(mask)
                    s_weight = nd.where(
                        mask > 0,
                        ones * self.LP_positive_weight,
                        ones * self.LP_negative_weight,
                        ctx=self.ctx[gpu_i])

                loss = self._get_loss_LP(by_, by, s_weight, mask)

                # ---------- Backward ---------- #
                sum(loss).backward()

        # ---------- Update Weights ---------- #
        self.trainer.step(self.batch_size)
        self.backward_counter += 1  # do not save at first step

        # ---------- Record Loss ---------- #
        if self.record and self.backward_counter % 10 == 0:
            # only record last gpu loss
            yolo_gluon.record_loss(loss, self.loss_name, self.sw,
                                   step=self.backward_counter,
                                   exp=self.exp)

        # ---------- Save Weights ---------- #
        if self.backward_counter % self.record_step == 0:
            idx = self.backward_counter // self.record_step
            path = os.path.join(self.backup_dir, self.exp + '_%d' % idx)
            self.net.collect_params().save(path)

    def _get_loss_LP(self, by_, by, s_weight, mask):
        s = self.LG_loss(by_[0], by[0], s_weight * self.scale['LP_score'])
        xy = self.HB_loss(by_[1], by[1], mask * self.scale['LP_xy'])
        z = self.HB_loss(by_[2], by[2], mask * self.scale['LP_z'])
        r = self.HB_loss(by_[3], by[3], mask * self.scale['LP_r'])
        c = self.CE_loss(by_[4], by[4], mask * self.scale['LP_class'])
        return (s, xy, z, r, c)


if __name__ == '__main__':
    main()