inference3

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

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'

startFrame = 1000

# windows
original_name = "Original"
result_name = "Result"
#win_size = cv2.WINDOW_AUTOSIZE

cv2.namedWindow(original_name, cv2.WINDOW_NORMAL)
cv2.namedWindow(result_name, cv2.WINDOW_NORMAL)
#cv2.setWindowProperty(original_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)

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'])
    parser.add_argument("--start-frame", type=str, default='0',
                        help="The frame of video to start infering.")

    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

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']

#frame = None
frame = cv2.imread(args.img_path)

frameId = 0
result = None

# 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)

# 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)

# Load the network
print("[INFO] loading 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)

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.')

class MyThread(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)

    def run(self):
        global result
        global frame
        global frameId

        while (~(frameId == 0)):
            result = self.predict(frame, frameId)

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

        preds = sess.run(pred, feed_dict={image_filename: name})
        return preds[0]


# 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")

cap.set(cv2.CAP_PROP_POS_FRAMES, startFrame)

# bool flag to play video
playVideo = True

tf_thread = MyThread()
tf_thread.start()

# Read until video is completed
while (cap.isOpened()):

    if playVideo:
        # Capture frame-by-frame
        ret, frame = cap.read()
        frameId = int(cap.get(cv2.CAP_PROP_POS_FRAMES)) # current frame number

        if ret == True:
            # Display the resulting frame
            cv2.imshow(original_name, frame)
            cv2.imshow(result_name, result)

        else:
            break

    # Press ESC on keyboard to exit
    key = cv2.waitKey(1) & 0xFF
    if key == ord('q'):
        break
    elif key == ord('p'):
        playVideo = not playVideo
    elif key == ord('f'):
        cv2.setWindowProperty(original_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
        cv2.setWindowProperty(result_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
    elif key == ord('a'):

        cv2.setWindowProperty(original_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_NORMAL)
        cv2.setWindowProperty(result_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_NORMAL)

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

sess.close()

# Closes all the frames
cv2.destroyAllWindows()