train.py

#-*- coding:utf-8 -*-
import os
import argparse
import numpy as np
import time
import sys
import argparse
import paddle
import paddle.fluid as fluid
import cv2
from collections import Counter

from models.unet import unet
from models.PAN import pannet
from models.dense_unet import dense_unet
from models.multires_unet import multires_unet
from models.deeplabv3p import deeplab_v3p
from models.deeplabv3p_ours import deeplabv3p_ours

from reader import TrainDataReader
import reader

os.environ['FLAGS_fraction_of_gpu_memory_to_use'] = '0.99'


pretrain_model = 0 #1 means pretrain_model
total_step = 150000
path = os.getcwd()

def saveImage(img,path):

    label_rgb = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)
    label_rgb[img == 0] =  [  0,   0,   0]
    label_rgb[img == 1] =  [ 70, 130, 180]
    label_rgb[img == 2] =  [119,  11,  32]
    label_rgb[img == 3] =  [220, 220,   0]
    label_rgb[img == 4] =  [102, 102, 156]
    label_rgb[img == 5] =  [128,  64, 128]
    label_rgb[img == 6] =  [190, 153, 153]
    label_rgb[img == 7] =  [128, 128,   0]
    label_rgb[img == 8] =  [255, 128,   0]
    img = cv2.resize(label_rgb, (4000,4000), interpolation=cv2.INTER_NEAREST)
    cv2.imwrite(path, label_rgb)

def categorical_crossentropy(label, pred):
    pred = fluid.layers.clip(x=pred, min=1e-7, max=1. - 1e-7)
    label = fluid.layers.reshape(label, shape=[-1,9])
    print('label.shape-----||||------pred.shape',label.shape,pred.shape)
    loss = fluid.layers.cross_entropy(input=pred, label=label,soft_label=True)
    #loss = fluid.layers.softmax_with_cross_entropy(logits=pred,label=label,ignore_index=0)
    loss = fluid.layers.reduce_mean(loss)
    loss = fluid.layers.clip(x=loss, min=3e-7, max=3. - 1e-7)
    return loss

def focal_loss(target,output):
    output = fluid.layers.clip(x=output, min=1e-7, max=1. - 1e-7)
    target = fluid.layers.reshape(target, shape=[-1,9])
    log = fluid.layers.log(output)
    log = fluid.layers.abs(log)
    log = fluid.layers.elementwise_mul(target,log)
    one = fluid.layers.fill_constant_batch_size_like(input=output, shape=output.shape, value=1.0, dtype='float32')
    one_sub_pt = fluid.layers.elementwise_sub(one,output)
    one_sub_pt_pow = fluid.layers.elementwise_mul(one_sub_pt,one_sub_pt)
    totalCost = fluid.layers.elementwise_mul(one_sub_pt_pow,log)
    loss = fluid.layers.reduce_mean(totalCost)
    loss = fluid.layers.clip(x=loss, min=1e-7, max=1. - 1e-7)
    return loss

def create_iou(predict, label, num_classes):
    print('iou_label.shape',predict.shape,label.shape)
    label = fluid.layers.reshape(label, shape=[-1,9])
    print('iou_label.shape',predict.shape,label.shape)
    label = fluid.layers.argmax(label, axis=1).astype('int64')
    predict = fluid.layers.argmax(predict, axis=1).astype('int64')
    print('iou_label.shape',predict.shape,label.shape)
    iou, out_wrong, out_right = fluid.layers.mean_iou(predict, label, num_classes)

    return iou,out_wrong,out_right

def cal_mean_iou(wrong, correct):
    sum = wrong + correct
    true_num = (sum != 0).sum()
    for i in range(len(sum)):
        if sum[i] == 0:
            sum[i] = 1
    return (correct.astype("float64") / sum).sum() / true_num

def create_reader(rows=1024,cols=1024):

    LaneDataset = reader.TrainDataReader
    #dataset = LaneDataset("/media/airobot/docs/BaiduDatas/ApolloDatas/", 'train',
    #                        rows=1024, cols=1024)
    dataset = LaneDataset(path + "/data/ApolloDatas/", 'train',
                            rows=1024, cols=1024)
    return dataset

def create_model(model='',image_shape=[1024,1024],class_num=9):

    train_image = fluid.layers.data(name='img', shape=[3] + image_shape, dtype='float32')
    train_label = fluid.layers.data(name='label', shape=image_shape + [9],dtype='float32')

    if model == 'unet':
        predict = unet().model(train_image)
    if model == 'deeplab_v3p':
        predict = deeplab_v3p().model(train_image)
    if model == 'pannet':
        predict = pannet().model(train_image)
    if model == 'dense_unet':
        predict = dense_unet().model(train_image)
    if model == 'multires_unet':
        predict = multires_unet().model(train_image)
    if model == 'deeplabv3p_ours':
        predict = deeplabv3p_ours().model(train_image)

    loss = categorical_crossentropy(train_label,predict)
    iou,out_wrong,out_right = create_iou(predict,train_label,class_num)

    return predict,loss,iou

def load_model(exe,program,model=''):
    if model == 'unet':
        fluid.io.load_params(executor=exe, dirname="", filename=path+'/params/unet.params', main_program=program)
    if model == 'deeplab_v3p':
        fluid.io.load_params(executor=exe, dirname="", filename=path+'/params/deeplab_v3p.params', main_program=program)
    if model == 'pannet':
        fluid.io.load_params(executor=exe, dirname="", filename=path+'/params/pannet.params', main_program=program)
    if model == 'dense_unet':
        fluid.io.load_params(executor=exe, dirname="", filename=path+'/params/dense_unet.params', main_program=program)
    if model == 'multires_unet':
        fluid.io.load_params(executor=exe, dirname="", filename=path+'/params/multires_unet.params', main_program=program)
    if model == 'deeplabv3p_ours':
        fluid.io.load_params(executor=exe, dirname="", filename=path+'/params/deeplabv3p_ours.params', main_program=program)


def save_model(exe,program,model=''):
    if model == 'unet':
        fluid.io.save_params(executor=exe, dirname="", filename=path+'/params/unet.params', main_program=program)
    if model == 'deeplab_v3p':
        fluid.io.save_params(executor=exe, dirname="", filename=path+'/params/deeplab_v3p.params', main_program=program)
    if model == 'pannet':
        fluid.io.save_params(executor=exe, dirname="", filename=path+'/params/pannet.params', main_program=program)
    if model == 'dense_unet':
        fluid.io.save_params(executor=exe, dirname="", filename=path+'/params/dense_unet.params', main_program=program)
    if model == 'multires_unet':
        fluid.io.loadsave_params_params(executor=exe, dirname="", filename=path+'/params/multires_unet.params', main_program=program)
    if model == 'deeplabv3p_ours':
        fluid.io.loadsave_params_params(executor=exe, dirname="", filename=path+'/params/deeplabv3p_ours.params', main_program=program)




def train(model):

    predict,loss,iou = create_model(model=model)
    optimizer = fluid.optimizer.Adam(learning_rate=1e-4)
    optimizer.minimize(loss)
    place = fluid.CUDAPlace(0)
    exe = fluid.Executor(place)

    exe.run(fluid.default_startup_program())
    fluid.memory_optimize(fluid.default_main_program(),print_log=False, skip_opt_set=set([loss.name,predict.name]))

    if pretrain_model:
        load_model(exe,fluid.default_main_program(),model=model)
        print("load model succeed")
    else:
        print("load succeed")

    def trainLoop():
        batches = DataSet.get_batch_generator(1, total_step)
        iou_count = 0
        mean_iou  = 0
        iou_sum   = 0
        for i, imgs, labels, names in batches:
            preTime = time.time()
            result = exe.run(fluid.default_main_program(),
                            feed={'img': imgs,
                                    'label': labels},
                           fetch_list=[loss,predict,iou])
            nowTime = time.time()

            iou_sum   += result[2]
            iou_count += 1
            mean_iou = iou_sum/iou_count

            print('                                                         iou = ',result[2],'mean_iou = ',mean_iou)
            
            if iou_count % 1000 == 0:
                iou_count = 0
                iou_sum   = 0

            if i % 1000 == 0 and i!= 0:
                print("Model saved")
                save_model(exe,fluid.default_main_program(),model=model)
            
            if i % 10 ==0:
                train_path = path+'/train.png'
                picture = result[1]
                picture = np.argmax(picture,axis=-1)
                picture = picture.reshape((1024,1024))
                saveImage(picture,train_path)
                label_path = path+'/trainlabel.png'
                train_lab = np.argmax(labels[0],axis=2)
                saveImage(train_lab,label_path)
                
            if i % 20 == 0:
                argmax = np.argmax(result[1],axis=1)
                abc = Counter(argmax)
                print('                                        ',abc)

            if i % 2 == 0:
                print("step {:d},loss {:.6f},step_time: {:.3f}".format(
                    i,result[0][0],nowTime - preTime))

    trainLoop()


if __name__ == "__main__":
    parse = argparse.ArgumentParser(description='')
    parse.add_argument('--model', help='model name', nargs='?')
    args = parse.parse_args()
    model = args.model
    DataSet = create_reader(model)
    train(model)