Add flask inference server and update readme with chrome extension and react app

inference_flask.py

@@ -0,0 +1,385 @@
+"""
+Copyright (C) 2020 NVIDIA Corporation.  All rights reserved.
+Licensed under the NVIDIA Source Code License. See LICENSE at https://github.com/nv-tlabs/GameGAN_code.
+Authors: Seung Wook Kim, Yuhao Zhou, Jonah Philion, Antonio Torralba, Sanja Fidler
+"""
+import os
+import sys
+import torch
+sys.path.append('..')
+import config
+import utils
+import random
+sys.path.insert(0, './data')
+import dataloader
+import copy
+import cv2
+import numpy as np
+import keyboard
+import time
+import base64
+
+import flask
+from flask import Flask, request, Response
+from flask import request
+from flask_cors import CORS,cross_origin
+
+import threading
+import json
+
+# Opencv AWS
+#sudo yum install mesa-libGL -y
+
+
+
+frames_dir = os.environ.get('SESSION_FRAMES_DIR', './session_frames') 
+if not os.path.exists(os.path.dirname(frames_dir)):
+    os.makedirs(os.path.dirname(frames_dir))
+
+# Workaround for PyTorch issue on Windows
+if os.name == 'nt':
+    import ctypes
+    ctypes.cdll.LoadLibrary('caffe2_nvrtc.dll')
+
+
+parser = config.init_parser()
+opts, args = parser.parse_args(sys.argv)
+gpu = opts.gpu
+
+# Initialize values
+opts = copy.deepcopy(opts)
+opts.img_size = (opts.img_size, opts.img_size)
+warm_up = opts.warm_up
+opts.gpu = gpu
+if opts.data is not None:
+    opts.num_data_types = len(opts.data.split('-'))
+
+log_dir = opts.log_dir
+
+# Multi-part model?
+if not os.path.exists(opts.saved_model):
+    part_list = sorted([file for file in os.listdir(os.path.dirname(opts.saved_model)) if os.path.basename(opts.saved_model) in file])
+    if len(part_list):
+        with open(opts.saved_model, 'wb') as sf:
+            for part in part_list:
+                with open(os.path.dirname(opts.saved_model) + '/' + part, 'rb') as lf:
+                    sf.write(lf.read())
+
+# Load the model
+saved_model = torch.load(opts.saved_model, map_location='cpu')
+opts_data = opts.data
+opts_img = opts.inference_image_path
+base_imgs = opts.show_base_images
+upsample_model = opts.upsample_model
+playback_fps = opts.playback_fps
+opts = saved_model['opts']
+if opts_data is not None:
+    opts.data = opts_data
+#if opts_img is not None:
+opts.inference_image_path = opts_img
+opts.show_base_images = base_imgs
+opts.upsample_model = upsample_model
+opts.playback_fps = playback_fps
+opts.gpu = gpu
+if type(opts.img_size) == int:
+    opts.img_size = [opts.img_size] * 2
+opts.log_dir = log_dir
+warm_up = opts.warm_up
+
+
+curdata, datadir = opts.data.split(':')
+if curdata == 'cartpole':
+    resized_image_size = (600, 400)
+    action_left = [1, 0]
+    action_right = [0, 1]
+    no_action = None
+elif curdata == 'vroom':
+    resized_image_size = (256, 256)
+    action_left = [1, 0, 0]
+    action_right = [0, 0, 1]
+    no_action = [0, 1, 0]
+elif curdata == 'gtav':
+    resized_image_size = (320, 192)
+    action_left = [1, 0, 0]
+    action_right = [0, 0, 1]
+    no_action = [0, 1, 0]
+else:
+    raise Exception(f'Not implemented: unknown data type: {curdata}')
+
+# Initialize torch
+torch.manual_seed(opts.seed)
+torch.cuda.set_device(gpu)
+
+# Create the generator model and load it;s state from the checkpoint
+netG, _ = utils.build_models(opts)
+utils.load_my_state_dict(netG, saved_model['netG'])
+
+# Initialize the noise generator
+zdist = utils.get_zdist('gaussian', opts.z)
+
+# For "playing", we want teh batch size of 1
+opts.bs = 1
+batch_size = opts.bs
+
+
+if opts.inference_image_path is None:
+    # Load the dataset so we can get some initial image
+    ##!! Replace with some examle set-aside images
+    train_loader = dataloader.get_custom_dataset(opts, set_type=0, getLoader=True)
+    data_iters, train_len = [], 99999999999
+    data_iters.append(iter(train_loader))
+    if len(data_iters[-1]) < train_len:
+        train_len = len(data_iters[-1])
+    states, actions, _ = utils.get_data(data_iters, opts)
+else:
+    # Load starting image
+    img = cv2.imread(opts.inference_image_path)[...,::-1]
+    img = (np.transpose(img, axes=(2, 0, 1)) / 255.).astype('float32')
+    img = (img - 0.5) / 0.5
+
+    states = [torch.tensor([img], dtype=torch.float32).cuda()]
+    actions = [torch.tensor(no_action if no_action is not None else action_left, dtype=torch.float32).cuda()]
+
+# Disable gradients, we'll perform just the inference
+utils.toggle_grad(netG, False)
+netG.eval()
+
+# Disable warmup
+warm_up = 0
+
+# Temporary, to save predictions as videos
+#fourcc = cv2.VideoWriter_fourcc(*'MP4V')
+#v = cv2.VideoWriter('video.mp4', fourcc, 10.0, resized_image_size)
+
+upsample = None
+if opts.upsample_model != None:
+    from upsample import upsample
+    upsample.load(opts.upsample_model)
+
+
+def apply_user_action(session_id, user_action):
+    encoding = 'utf-8'
+
+    action = None
+    hidden_action = 0
+    prev_state = None
+    i = 0
+    M = None
+    prev_read_v = None
+    prev_alpha = None
+
+    torch_file = "%s/%s.pt" % (frames_dir, session_id)
+    torch_file_c = "%s/c_%s.pt" % (frames_dir, session_id)
+    torch_file_h = "%s/h_%s.pt" % (frames_dir, session_id)
+    other_vars_file = "%s/vars_%s.json" % (frames_dir, session_id)
+
+    if os.path.exists(torch_file):
+        prev_state = torch.load(torch_file)
+        h = [torch.load(torch_file_c)]
+        c = [torch.load(torch_file_h)]
+
+    if os.path.exists(other_vars_file):
+        with open(other_vars_file) as json_file:
+            other_vars_data = json.load(json_file)
+            i = other_vars_data["steps"]
+            M = other_vars_data["M"]
+            prev_read_v = other_vars_data["prev_read_v"]
+            prev_alpha = other_vars_data["prev_alpha"]
+    i += 1
+
+
+    frame_start_time = time.time()
+
+    action_text = ''
+    if user_action == 'r' or prev_state is None:
+        # Run warmup to get initial values
+        # warmup is set to 0, so initial image is going to be used as input
+        prev_state, warm_up_state, M, prev_read_v, prev_alpha, outputs, maps, alphas, alpha_losses, zs, base_imgs_all, _, \
+            hiddens, init_maps = netG.run_warmup(zdist, states, actions, warm_up, train=False)
+        h, c = warm_up_state
+
+        # Show the image
+        img = prev_state[0].cpu().numpy()
+        img = np.rollaxis(img, 0, 3)
+        img = ((img+1)*127.5).astype(np.uint8)
+        img = cv2.resize(img, resized_image_size, interpolation=cv2.INTER_NEAREST)
+        img = img[...,::-1]
+
+        _, im_arr = cv2.imencode('.png', img)  # im_arr: image in Numpy one-dim array format.
+        im_bytes = im_arr.tobytes()
+        inference_base64 = str(base64.b64encode(im_bytes), encoding)
+
+        #cv2.imshow(f'{curdata} - inference', img)
+        if upsample is not None:
+            upsampled_img = upsample.inference(np.rollaxis(prev_state[0].cpu().numpy(), 0, 3))
+
+            upsampled_base64 = None
+            if upsample is not None:
+                upsampled_img = upsample.inference(np.rollaxis(prev_state[0].cpu().numpy(), 0, 3))
+                if type(upsampled_img) is np.ndarray:
+                    _, im_arr = cv2.imencode('.png', upsampled_img[0][...,::-1])  # im_arr: image in Numpy one-dim array format.
+                    im_bytes = im_arr.tobytes()
+                    upsampled_base64 = str(base64.b64encode(im_bytes), encoding)
+                else:
+                    _, im_arr = cv2.imencode('.png', upsampled_img[0][0][...,::-1])  # im_arr: image in Numpy one-dim array format.
+                    im_bytes = im_arr.tobytes()
+                    upsampled_base64 = str(base64.b64encode(im_bytes), encoding)
+
+        result = {'inference': inference_base64}
+        if upsampled_base64 is not None:
+            result['upscaled'] = upsampled_base64
+
+
+        other_vars_data = {
+            "steps": i,
+            "M": M,
+            "prev_read_v": prev_read_v,
+            "prev_alpha": prev_alpha,
+        }
+
+        threading.Thread(target=save_torch_state, args=(prev_state, torch_file, h[0], torch_file_h, c[0], torch_file_c, other_vars_data, other_vars_file)).start()
+
+        return result
+
+    elif user_action == 'a':
+        action = torch.tensor([action_left], dtype=torch.float32).cuda()
+        hidden_action = -1
+        #action_text = 'LEFT'
+    elif user_action == 'd':
+        action = torch.tensor([action_right], dtype=torch.float32).cuda()
+        hidden_action = 1
+        #action_text = 'RIGHT'
+    elif no_action is not None:
+        action = torch.tensor([no_action], dtype=torch.float32).cuda()
+        hidden_action = 0
+    else:
+        action = torch.tensor([np.eye(opts.action_space)[random.randint(0, np.eye(opts.action_space) - 1)]], dtype=torch.float32).cuda()
+        hidden_action = None
+
+    # Perform inference
+    prev_state, m, prev_alpha, alpha_loss, z, M, prev_read_v, h, c, init_map, base_imgs, _, cur_hidden = netG.run_step(prev_state, h, c, action, \
+                                                                          batch_size, prev_read_v, prev_alpha, M, zdist, step=i)
+
+    # Show the image
+    img = prev_state[0].cpu().numpy()
+    img = np.rollaxis(img, 0, 3)
+    img = ((img+1)*127.5).astype(np.uint8)
+    img = cv2.resize(img, resized_image_size, interpolation=cv2.INTER_NEAREST)
+    img = img[...,::-1]
+    rectangle = img.copy()
+    cv2.rectangle(rectangle, (0, 0), (150, 30), (0, 0, 0), -1)
+    img = cv2.addWeighted(rectangle, 0.6, img, 0.4, 0)
+    cv2.putText(img, "Action:", (8, 18), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
+    if hidden_action == -1:
+        color = (55, 155, 255)
+        text = "LEFT"
+    elif hidden_action == 1:
+        text = "RIGHT"
+        color = (55, 155, 255)
+    elif hidden_action == 0:
+        text = "STRAIGHT"
+        color = (55, 255, 55)
+    else:
+        text = "UNKNOWN"
+        color = (55, 55, 255)
+    cv2.putText(img, text, (80, 18), cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
+    _, im_arr = cv2.imencode('.png', img)  # im_arr: image in Numpy one-dim array format.
+    im_bytes = im_arr.tobytes()
+    inference_base64 = str(base64.b64encode(im_bytes), encoding)
+
+    if opts.show_base_images == 'True' and opts.num_components > 1:
+        for i in range(opts.num_components * 2):
+            img = base_imgs[i][0].cpu().numpy()
+            img = np.rollaxis(img, 0, 3)
+            img = ((img+1)*127.5).astype(np.uint8)
+            img = cv2.resize(img, resized_image_size, interpolation=cv2.INTER_NEAREST)
+            img = img[...,::-1]
+        for i in range(opts.num_components):
+            img = m[i][0].cpu().numpy()
+            img = np.rollaxis(img, 0, 3)
+            img = ((img+1)*127.5).astype(np.uint8)
+            img = np.expand_dims(cv2.resize(img, resized_image_size, interpolation=cv2.INTER_NEAREST), -1)
+            img = img[...,::-1]
+
+    upsampled_base64 = None
+    if upsample is not None:
+        upsampled_img = upsample.inference(np.rollaxis(prev_state[0].cpu().numpy(), 0, 3))
+        if type(upsampled_img) is np.ndarray:
+            _, im_arr = cv2.imencode('.png', upsampled_img[0][...,::-1])  # im_arr: image in Numpy one-dim array format.
+            im_bytes = im_arr.tobytes()
+            upsampled_base64 = str(base64.b64encode(im_bytes), encoding)
+        else:
+            _, im_arr = cv2.imencode('.png', upsampled_img[0][0][...,::-1])  # im_arr: image in Numpy one-dim array format.
+            im_bytes = im_arr.tobytes()
+            upsampled_base64 = str(base64.b64encode(im_bytes), encoding)
+
+    result = {'inference': inference_base64}
+    if upsampled_base64 is not None:
+        result['upscaled'] = upsampled_base64
+
+    other_vars_data = {
+        "steps": i,
+        "M": m,
+        "prev_read_v": prev_read_v,
+        "prev_alpha": prev_alpha,
+    }
+
+    threading.Thread(target=save_torch_state, args=(prev_state, torch_file, h[0], torch_file_h, c[0], torch_file_c, other_vars_data, other_vars_file)).start()
+
+    return result
+
+def save_torch_state(prev_state, torch_file, h, torch_file_h, c, torch_file_c, other_vars_data, other_vars_file):
+    torch.save(prev_state, torch_file)
+    torch.save(h, torch_file_h)
+    torch.save(c, torch_file_c)
+
+    with open(other_vars_file, 'w', encoding='utf-8') as f:
+        json.dump(other_vars_data, f, ensure_ascii=False, indent=4)
+        #    with open(other_vars_file, 'w') as f:
+        #json.dump(other_vars_data, f)
+
+app = Flask(__name__)
+CORS(app)
+app.config['CORS_HEADERS'] = 'Content-Type'
+
+DEBUG = bool(int(os.environ.get('DEBUG', 1)))
+FLASK_PORT = int(os.environ.get('FLASK_PORT', 8751))
+
+@app.route('/api/sample_step',methods=['GET', 'POST', 'OPTIONS'])
+@cross_origin(origin='*',headers=['access-control-allow-origin','Content-Type'])
+@cross_origin()
+def sample_step():
+    payload = {}
+
+    params = request.get_json()
+    session_id = params['session_id']
+
+    result = apply_user_action(session_id, random.choice(['a', 'd']))
+
+    return json.dumps(result)
+
+@app.route('/api/step',methods=['GET', 'POST', 'OPTIONS'])
+@cross_origin(origin='*',headers=['access-control-allow-origin','Content-Type'])
+@cross_origin()
+def step():
+    payload = {}
+
+    params = request.get_json()
+    session_id = params['session_id']
+    user_action = params['action']
+
+    result = apply_user_action(session_id, user_action)
+
+    return json.dumps(result)
+
+@app.route('/api/healthcheck', methods=['GET', 'OPTIONS'])
+@cross_origin(origin='*',headers=['access-control-allow-origin','Content-Type'])
+@cross_origin()
+def health_check():
+
+    return json.dumps({})
+
+
+
+if __name__ == '__main__':
+    app.run(debug=DEBUG,host='0.0.0.0', port=FLASK_PORT, threaded=False)