test-multithreading

from __future__ import print_function

import argparse
import os
import sys
import time
import tensorflow as tf
import numpy as np
from scipy import misc
import cv2
import multiprocessing
from matplotlib import pyplot as plt
from imutils.video import FPS
from multiprocessing import Queue, Pool

from model import PSPNet101, PSPNet50
from tools import *

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

ADE20k_param = {'crop_size': [473, 473],
                'num_classes': 150,
                'model': PSPNet50}
cityscapes_param = {'crop_size': [720, 720],
                    'num_classes': 19,
                    'model': PSPNet101}

SAVE_DIR = './output/'
SNAPSHOT_DIR = './model/'
VIDEO_PATH = './input/085853AA.mp4'

def get_arguments():
    parser = argparse.ArgumentParser(description="Reproduced PSPNet")
    parser.add_argument("--img-path", type=str, default='',
                        help="Path to the RGB image file.")
    parser.add_argument("--checkpoints", type=str, default=SNAPSHOT_DIR,
                        help="Path to restore weights.")
    parser.add_argument("--save-dir", type=str, default=SAVE_DIR,
                        help="Path to save output.")
    parser.add_argument("--flipped-eval", action="store_true",
                        help="whether to evaluate with flipped img.")
    parser.add_argument("--dataset", type=str, default='cityscapes',
                        choices=['ade20k', 'cityscapes'])
    parser.add_argument("--source", type=str, default='video',
                        choices=['camera', 'video'])

    return parser.parse_args()

def save(saver, sess, logdir, step):
   model_name = 'model.ckpt'
   checkpoint_path = os.path.join(logdir, model_name)

   if not os.path.exists(logdir):
      os.makedirs(logdir)
   saver.save(sess, checkpoint_path, global_step=step)
   print('The checkpoint has been created.')

def load(saver, sess, ckpt_path):
    saver.restore(sess, ckpt_path)
    print("Restored model parameters from {}".format(ckpt_path))

def create_blank(height, width, rgb_color=(0, 0, 0)):
    """Create new image(numpy array) filled with certain color in RGB"""
    # Create black blank image
    image = np.zeros((height, width, 3), np.uint8)

    # Since OpenCV uses BGR, convert the color first
    color = tuple(reversed(rgb_color))
    # Fill image with color
    image[:] = color

    return image

def twoImgShowVertically(img1, img2, border):
    h, w, channels = img1.shape

    img2 = cv2.resize(img2, (w, h))
    both = create_blank(h + border * 2, w * 2 + border * 3)

    both[border:border + h, border:border + w] = img1.copy()
    both[border:border + h, border * 2 + w:border * 2 + w * 2] = img2.copy()

    return both

def twoImgShowHorizontally(img1, img2, border):
    h, w, channels = img1.shape

    img2 = cv2.resize(img2, (w, h))
    both = create_blank(h * 2 + border * 3, w + border * 2)

    both[border:border + h, border:border + w] = img1.copy()
    both[border * 2 + h:border * 2 + h * 2, border:border + w] = img2.copy()

    return both

def worker(input_q, output_q):
    args = get_arguments()

    param = cityscapes_param

    crop_size = param['crop_size']
    num_classes = param['num_classes']
    PSPNet = param['model']

    # preprocess images
    image_filename = tf.placeholder(dtype=tf.string)
    img = tf.image.decode_image(tf.read_file(image_filename), channels=3)
    img.set_shape([None, None, 3])

    img_shape = tf.shape(img)
    h, w = (tf.maximum(crop_size[0], img_shape[0]), tf.maximum(crop_size[1], img_shape[1]))
    img = preprocess(img, h, w)

    # Create network.
    net = PSPNet({'data': img}, is_training=False, num_classes=num_classes)
    with tf.variable_scope('', reuse=True):
        flipped_img = tf.image.flip_left_right(tf.squeeze(img))
        flipped_img = tf.expand_dims(flipped_img, dim=0)
        net2 = PSPNet({'data': flipped_img}, is_training=False, num_classes=num_classes)

    raw_output = net.layers['conv6']

    # Do flipped eval or not
    if args.flipped_eval:
        flipped_output = tf.image.flip_left_right(tf.squeeze(net2.layers['conv6']))
        flipped_output = tf.expand_dims(flipped_output, dim=0)
        raw_output = tf.add_n([raw_output, flipped_output])

    # Predictions.
    raw_output_up = tf.image.resize_bilinear(raw_output, size=[h, w], align_corners=True)
    raw_output_up = tf.image.crop_to_bounding_box(raw_output_up, 0, 0, img_shape[0], img_shape[1])
    raw_output_up = tf.argmax(raw_output_up, axis=3)
    pred = decode_labels(raw_output_up, img_shape, num_classes)

    # Init tf Session
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
    sess = tf.Session(config=config)
    init = tf.global_variables_initializer()

    sess.run(init)

    restore_var = tf.global_variables()

    # Load checkpoint
    ckpt = tf.train.get_checkpoint_state(args.checkpoints)
    if ckpt and ckpt.model_checkpoint_path:
        loader = tf.train.Saver(var_list=restore_var)
        load_step = int(os.path.basename(ckpt.model_checkpoint_path).split('-')[1])
        load(loader, sess, ckpt.model_checkpoint_path)
    else:
        print('No checkpoint file found.')

    while True:
        name = input_q.get()
        preds = sess.run(pred, feed_dict={image_filename: name})
        output_q.put(preds[0])

    sess.close()

def main():
    args = get_arguments()

    # Capture the video source
    if args.source == 'camera':
        cap = cv2.VideoCapture(0)
    elif args.source == 'video':
        cap = cv2.VideoCapture(VIDEO_PATH)

    logger = multiprocessing.log_to_stderr()
    logger.setLevel(multiprocessing.SUBDEBUG)

    input_q = Queue(maxsize=5)
    output_q = Queue(maxsize=5)
    pool = Pool(4, worker, (input_q, output_q))

    # Check if camera opened successfully
    if (cap.isOpened() == False):
        print("Error opening video stream or file")

    fps = FPS().start()

    # Read until video is completed
    while(cap.isOpened()):
        # Capture frame-by-frame
        ret, frame = cap.read()
        frameId = cap.get(1) # current frame number
        frame = cv2.resize(frame, None, fx=1/6, fy=1/6)

        name = './tmp/' + str(frameId) + '.jpg'
        cv2.imwrite(name, frame)

        if ret == True:
            input_q.put(name)

            cv2.imshow('Original', frame)
            cv2.imshow('Result', output_q.get())

            # Press ESC on keyboard to  exit
            if cv2.waitKey(33) & 0xFF == ord('q'):
                break

            # update the FPS counter
            fps.update()

        else:
            break

        os.remove(name)

    # stop the timer and display FPS information
    fps.stop()
    print("[INFO] elapsed time: {:.2f}".format(fps.elapsed()))
    print("[INFO] approx. FPS: {:.2f}".format(fps.fps()))

    # When everything done, release the video capture object
    cap.release()

    # Closes all the frames
    cv2.destroyAllWindows()

if __name__ == '__main__':
    main()