test-multithreading2

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
from matplotlib import pyplot as plt

from imutils.video import VideoStream
from imutils.video import FPS
from multiprocessing import Process
from multiprocessing import Queue
import imutils

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 main():
    args = get_arguments()

    ###########################################################
    # load parameters
    if args.dataset == 'ade20k':
        param = ADE20k_param
    elif args.dataset == 'cityscapes':
        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.
    print("[INFO] loading model...")
    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()
    #print(restore_var)
    # 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.')

    def classify_frame(model, inputQueue, outputQueue):
        while True:
            if not inputQueue.empty():
                # grab the frame from the input queue, resize it, and
                # construct a blob from it
                input_name = inputQueue.get()
                print(input_name)

                preds = sess.run(model, feed_dict={image_filename: input_name})
                prediction = preds[0]

                # write the detections to the output queue
                outputQueue.put(prediction)

    # initialize the input queue (detected faces grayed), output queue (predictions),
    inputQueue = Queue(maxsize=4)
    outputQueue = Queue(maxsize=4)
    out = None
    #cv2.imshow('abco', out)

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

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

    # initialize the prediction process
    p = Process(target=classify_frame, args=(pred, inputQueue, outputQueue,))
    p.daemon = True
    p.start()

    fps = FPS().start()

    # Read until video is completed
    while(cap.isOpened()):
        # Capture frame-by-frame
        ret, frame = cap.read()
        frameId = int(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:

            # if the input queue *is* empty, give the current frame to
            # classify
            if inputQueue.empty():
                inputQueue.put(name)

            # if the output queue *is not* empty, grab the detections
            if not outputQueue.empty():
                out = outputQueue.get()
            #else:
            #    print('outputQueue is empty')

            cv2.imshow('Original', frame)
            cv2.imshow('Result', out)
            # 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()

    sess.close()

    # Closes all the frames
    cv2.destroyAllWindows()


if __name__ == '__main__':
    main()