train.py

import argparse, glob, os, cv2, sys, pickle, random
import numpy as np
import tensorflow as tf
import config as cfg
from models.stgru import STGRU
from models.lrr import LRR
from models.dilation import dilation10network
from models.flownet2 import Flownet2
from models.flownet1 import Flownet1
from tensorflow.python.framework import ops

bilinear_warping_module = tf.load_op_library('./misc/bilinear_warping.so')
@ops.RegisterGradient("BilinearWarping")
def _BilinearWarping(op, grad):
  return bilinear_warping_module.bilinear_warping_grad(grad, op.inputs[0], op.inputs[1])

class DataLoader():
    def __init__(self, im_size, nbr_frames):
        self.im_size = im_size
        self.dataset_size = [1024, 2048]
        self.nbr_frames = nbr_frames
        self.L = glob.glob(os.path.join(cfg.cityscapes_dir, 'gtFine', 'train', "*", "*labelTrainIds.png"))
        random.shuffle(self.L)
        self.idx = 0

    def get_next_sequence(self):
        H, W = self.dataset_size
        h, w = self.im_size

        offset = [np.random.randint(H - h),
            np.random.randint(W - w)]
        i0, j0 = offset
        i1, j1 = i0 + h, j0 + w

        im_path = self.L[self.idx % len(self.L)]
        self.idx += 1

        parts = im_path.split('/')[-1].split('_')
        city, seq, frame = parts[0], parts[1], parts[2]

        images = []
        gt = cv2.imread(im_path, 0)[i0:i1, j0:j1]
        
        for dt in range(-self.nbr_frames + 1, 1):
            t = int(frame) + dt
            
            frame_path = os.path.join(cfg.cityscapes_video_dir, 'leftImg8bit_sequence', 'train', 
                    city, ("%s_%s_%06d_leftImg8bit.png" % (city, seq, t)))
            images.append(cv2.imread(frame_path, 1).astype(np.float32)[i0:i1,j0:j1][np.newaxis,...])
            
        return images, gt

def train(args):
    nbr_classes = 19

    # learning rates for the GRU and the static segmentation networks, respectively
    learning_rate = 2e-5
    static_learning_rate = 2e-12
    
    # The total number of iterations and when the static network should start being refined
    nbr_iterations = 10000
    t0_dilation_net = 5000

    im_size = [512, 512]
    image_mean = [72.39,82.91,73.16] # the mean is automatically subtracted in some modules e.g. flownet2, so be careful

    f = open('misc/cityscapes_labels.pckl')
    cs_id2trainid, cs_id2name = pickle.load(f)
    f.close()

    assert args.static in ['dilation', 'lrr'], "Only dilation and LRR are supported for now."

    if args.flow == 'flownet2':
        with tf.variable_scope('flow'):
            flow_network = Flownet2(bilinear_warping_module)
            flow_img0 = tf.placeholder(tf.float32)
            flow_img1 = tf.placeholder(tf.float32)
            flow_tensor = flow_network(flow_img0, flow_img1, flip=True)
    elif args.flow == 'flownet1':
        with tf.variable_scope('flow'):
            flow_network = Flownet1()
            flow_img0 = tf.placeholder(tf.float32)
            flow_img1 = tf.placeholder(tf.float32)
            flow_tensor = flow_network.get_output_tensor(flow_img0, flow_img1, im_size)

    RNN = STGRU([nbr_classes, im_size[0], im_size[1]], [7, 7], bilinear_warping_module)
    
    gru_opt, gru_loss, gru_prediction, gru_learning_rate, \
        gru_input_images_tensor, gru_input_flow_tensor, \
        gru_input_segmentation_tensor, gru_targets = RNN.get_optimizer(args.frames)
    unary_grad_op = tf.gradients(gru_loss, gru_input_segmentation_tensor)

    if args.static == 'lrr':
        static_input = tf.placeholder(tf.float32)
        static_network = LRR()
        static_output = static_network(static_input)

        unary_opt, unary_dLdy = static_network.get_optimizer(static_input, static_output, static_learning_rate)
    elif args.static == 'dilation':
        static_input = tf.placeholder(tf.float32)
        static_network = dilation10network()
        static_output = static_network.get_output_tensor(static_input, im_size)

    data_loader = DataLoader(im_size, args.frames)

    loss_history = np.zeros(nbr_iterations)
    loss_history_smoothed = np.zeros(nbr_iterations)

    vars_trainable = [k for k in tf.trainable_variables() if not k.name.startswith('flow/')]
    vars_static = [k for k in vars_trainable if not k in RNN.weights.values()]
    loader_static = tf.train.Saver(vars_static)
    saver = tf.train.Saver(vars_trainable)
    
    if args.flow in ['flownet1', 'flownet2']:
        saver_fn = tf.train.Saver([k for k in tf.trainable_variables() if k.name.startswith('flow/')])

    init = tf.global_variables_initializer()

    with tf.Session() as sess:
        sess.run(init)

        if args.static == 'lrr':
            loader_static.restore(sess, './checkpoints/lrr_pretrained')
        elif args.static == 'dilation':
            assert False, "Pretrained dilation model will soon be released."
            saver.restore(sess, './checkpoints/dilation_grfp')

        if args.flow == 'flownet1':
            saver_fn.restore(sess, './checkpoints/flownet1')
        elif args.flow == 'flownet2':
            saver_fn.restore(sess, './checkpoints/flownet2')

        for training_it in range(nbr_iterations):
            images, ground_truth = data_loader.get_next_sequence()

            # Optical flow
            optflow = []
            for frame in range(1, args.frames):
                im, last_im = images[frame], images[frame-1]
                if args.flow == 'flownet2':
                    flow = sess.run(flow_tensor, feed_dict={flow_img0: im, flow_img1: last_im})
                elif args.flow == 'flownet1':
                    flow = sess.run(flow_tensor, feed_dict={flow_img0: im, flow_img1: last_im})
                    flow = flow[...,(1, 0)]
                elif args.flow == 'farneback':
                    im_gray = cv2.cvtColor(im[0], cv2.COLOR_BGR2GRAY)
                    last_im_gray = cv2.cvtColor(last_im[0], cv2.COLOR_BGR2GRAY)

                    flow = cv2.calcOpticalFlowFarneback(im_gray, last_im_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0)
                    flow = flow[...,(1, 0)]
                    flow = flow[np.newaxis,...]
                optflow.append(flow)

            # Static segmentation
            static_segm = []
            for frame in range(args.frames):
                im = images[frame]
                if args.static == 'dilation':
                    # augment a 186x186 border around the image and subtract the mean
                    im_aug = cv2.copyMakeBorder(im[0], 186, 186, 186, 186, cv2.BORDER_REFLECT_101)
                    im_aug = im_aug - image_mean
                    im_aug = im_aug[np.newaxis,...]

                    x = sess.run(static_output, feed_dict={static_input: im_aug})
                elif args.static == 'lrr':
                    x = sess.run(static_output, feed_dict={static_input: im})
                static_segm.append(x)

            # GRFP
            rnn_input = {
                gru_learning_rate: learning_rate,
                gru_input_images_tensor: np.stack(images),
                gru_input_flow_tensor: np.stack(optflow),
                gru_input_segmentation_tensor: np.stack(static_segm),
                gru_targets: ground_truth,
            }

            _, loss, pred, unary_grads = sess.run([gru_opt, gru_loss, 
               gru_prediction, unary_grad_op], feed_dict=rnn_input)
            loss_history[training_it] = loss
            
            if training_it < 300:
                loss_history_smoothed[training_it] = np.mean(loss_history[0:training_it+1])
            else:
                loss_history_smoothed[training_it] = 0.997*loss_history_smoothed[training_it-1] + 0.003*loss

            # Refine the static network?
            # The reason that a two-stage training routine is used
            # is because there is not enough GPU memory (with a 12 GB Titan X)
            # to do it in one pass.
            if training_it+1 > t0_dilation_net:
                for k in range(len(images)-3, len(images)):
                    g = unary_grads[0][k]
                    im = images[k]
                    _ = sess.run([unary_opt], feed_dict={
                      static_input: im,
                      unary_dLdy: g
                    })

            if training_it > 0 and (training_it+1) % 1000 == 0:
                saver.save(sess, './checkpoints/%s_%s_it%d' % (args.static, args.flow, training_it+1))

            if (training_it+1) % 200 == 0:
                print("Iteration %d/%d: Loss %.3f" % (training_it+1, nbr_iterations, loss_history_smoothed[training_it]))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Tran GRFP on the CityScapes training set.')

    parser.add_argument('--static', help='Which static network to use.', required=True)
    parser.add_argument('--flow', help='Which optical flow method to use.', required=True)
    parser.add_argument('--frames', type=int, help='Number of frames to use.', default=5, required=False)

    args = parser.parse_args()

    assert args.flow in ['flownet1', 'flownet2', 'farneback'], "Unknown flow method %s." % args.flow
    assert args.static in ['dilation', 'dilation_grfp', 'lrr', 'lrr_grfp'], "Unknown static method %s." % args.static
    assert args.frames >= 1 and args.frames <= 20, "The number of frames must be between 1 and 20."
    
    train(args)