predict_masks.py

import os
from time import clock

import numpy as np
from keras.applications.imagenet_utils import preprocess_input
from keras.preprocessing.image import array_to_img, img_to_array, load_img, flip_axis

from models import make_model
from params import args

os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu

prediction_dir = args.pred_mask_dir


def do_tta(x, tta_type):
    if tta_type == 'hflip':
        # batch, img_col = 2
        return flip_axis(x, 2)
    else:
        return x


def undo_tta(pred, tta_type):
    if tta_type == 'hflip':
        # batch, img_col = 2
        return flip_axis(pred, 2)
    else:
        return pred


def predict():
    output_dir = args.pred_mask_dir
    model = make_model((None, None, 3))
    model.load_weights(args.weights)
    batch_size = args.pred_batch_size
    nbr_test_samples = 100064

    filenames = [os.path.join(args.test_data_dir, f) for f in sorted(os.listdir(args.test_data_dir))]

    start_time = clock()
    for i in range(int(nbr_test_samples / batch_size) + 1):
        x = []
        for j in range(batch_size):
            if i * batch_size + j < len(filenames):
                img = load_img(filenames[i * batch_size + j], target_size=(args.img_height, args.img_width))
                x.append(img_to_array(img))
        x = np.array(x)
        x = preprocess_input(x, args.preprocessing_function)
        x = do_tta(x, args.pred_tta)
        batch_x = np.zeros((x.shape[0], 1280, 1920, 3))
        batch_x[:, :, 1:-1, :] = x
        preds = model.predict_on_batch(batch_x)
        preds = undo_tta(preds, args.pred_tta)
        for j in range(batch_size):
            filename = filenames[i * batch_size + j]
            prediction = preds[j][:, 1:-1, :]
            array_to_img(prediction * 255).save(os.path.join(output_dir, filename.split('/')[-1][:-4] + ".png"))
        time_spent = clock() - start_time
        print("predicted batch ", str(i))
        print("Time spent: {:.2f}  seconds".format(time_spent))
        print("Speed: {:.2f}  ms per image".format(time_spent / (batch_size * (i + 1)) * 1000))
        print("Elapsed: {:.2f} hours  ".format(time_spent / (batch_size * (i + 1)) / 3600 * (nbr_test_samples - (batch_size * (i + 1)))))

if __name__ == '__main__':
    predict()