PREVIOUS_STATES_DISTANCES = [1, 2] MODEL_HEIGHT = 90 MODEL_WIDTH = 160 MODEL_PREV_HEIGHT = 45 MODEL_PREV_WIDTH = 80 class Predictor(nn.Module):
class PredictorWithShortcuts(nn.Module):
    def __init__(self):
        def identity(_):
    def downscale(self, img):
    def downscale_prev(self, img):
    def embed_state(self, previous_states, state, volatile=False, requires_grad=True, gpu=-1):
    def embed(self, inputs, inputs_prev):
    def forward(self, inputs, inputs_prev, only_embed=False):
        def act(x):
        def lrelu(x, negative_slope=0.2):
        def up(x, f=2):
        def maxp(x):
    def predict_grids(self, inputs, inputs_prev):
    def __init__(self):
        def identity(_):
    def downscale(self, img):
    def downscale_prev(self, img):
    def embed_state(self, previous_states, state, volatile=False, requires_grad=True, gpu=-1):
    def embed(self, inputs, inputs_prev):
    def forward(self, inputs, inputs_prev, only_embed=False):
        def act(x):
        def lrelu(x, negative_slope=0.2):
        def up(x, f=2):
        def maxp(x):
    def predict_grids(self, inputs, inputs_prev):
        B = inputs.size(0)
        B = inputs.size(0) def to_variable(inputs, volatile=False, requires_grad=True):
    if volatile:
        make_var = lambda x: Variable(x, volatile=True)
    else:
        make_var = lambda x: Variable(x, requires_grad=requires_grad)

    if isinstance(inputs, np.ndarray):
        return make_var(torch.from_numpy(inputs))
    elif isinstance(inputs, list):
        return [make_var(torch.from_numpy(el)) for el in inputs]
    elif isinstance(inputs, tuple):
        return [make_var(torch.from_numpy(el)) for el in inputs]
    elif isinstance(inputs, dict):
        return dict([(key, make_var(torch.from_numpy(inputs[key]))) for key in inputs])
    else:
        raise Exception("unknown input %s" % (type(inputs),)) def to_cuda(inputs, gpu=Config.GPU):
    if gpu <= -1:
        return inputs
    else:
        if isinstance(inputs, Variable):
            return inputs.cuda(gpu)
        elif isinstance(inputs, list):
            return [el.cuda(gpu) for el in inputs]
        elif isinstance(inputs, tuple):
            return tuple([el.cuda(gpu) for el in inputs])
        elif isinstance(inputs, dict):
            return dict([(key, inputs[key].cuda(gpu)) for key in inputs])
        else:
            raise Exception("unknown input %s" % (type(inputs),)) def to_numpy(var):
    #if ia.is_numpy_array(var):
    if isinstance(var, (np.ndarray, np.generic)):
        return var
    else:
        return var.data.cpu().numpy() def template_match(needle, haystack):
def compress_to_jpg(img):
def compress_img(img, method):
def decompress_img(img_compressed):
def draw_image(img, other_img, x, y, copy=True, raise_if_oob=False):
def draw_point(img, y, x, size=5, color=[255, 0, 0], alpha=1.0, copy=True):
def draw_line(img, y1, x1, y2, x2, color=[0, 255, 0], alpha=1.0, thickness=1, clip_result=True, copy=True):
def draw_direction_circle(img, y, x, r_inner, r_outer, angle_start, angle_end, color_border=[0, 255, 0], color_fill=[0, 200, 0], alpha=1.0):
def draw_text(img, y, x, text, color=[0, 255, 0], size=25):
def to_variable(inputs, volatile=False, requires_grad=True):
def to_variables(inputs, volatile=False, requires_grad=True):
def to_cuda(inputs, gpu=Config.GPU):
def to_numpy(var):
def draw_heatmap_overlay(img, heatmap, alpha=0.5):
def create_2d_gaussian_old(size, fwhm=3, center=None):
def create_2d_gaussian(size, sigma):
FONT_CACHE = dict() ACTION_UP = "W"
ACTION_DOWN = "S"
ACTION_LEFT = "A"
ACTION_RIGHT = "D"
ACTION_UP_DOWN_NONE = "~WS"
ACTION_LEFT_RIGHT_NONE = "~AD"
ACTIONS_UP_DOWN = [ACTION_UP, ACTION_DOWN, ACTION_UP_DOWN_NONE]
ACTIONS_LEFT_RIGHT = [ACTION_LEFT, ACTION_RIGHT, ACTION_LEFT_RIGHT_NONE]
ALL_ACTIONS = ACTIONS_UP_DOWN + ACTIONS_LEFT_RIGHT
ACTION_TO_KEY = {
    ACTION_UP: "w",
    ACTION_DOWN: "s",
    ACTION_LEFT: "a",
    ACTION_RIGHT: "d",
    ACTION_UP_DOWN_NONE: None,
    ACTION_LEFT_RIGHT_NONE: None
}
ACTION_UP_VECTOR_INDEX = 0
ACTION_DOWN_VECTOR_INDEX = 1
ACTION_LEFT_VECTOR_INDEX = 0
ACTION_RIGHT_VECTOR_INDEX = 1
ACTION_UP_DOWN_NONE_VECTOR_INDEX = 2
ACTION_LEFT_RIGHT_NONE_VECTOR_INDEX = 2
ACTION_UP_VECTOR = make_one_hot_vector(3, ACTION_UP_VECTOR_INDEX)
ACTION_DOWN_VECTOR = make_one_hot_vector(3, ACTION_DOWN_VECTOR_INDEX)
ACTION_LEFT_VECTOR = make_one_hot_vector(3, ACTION_LEFT_VECTOR_INDEX)
ACTION_RIGHT_VECTOR = make_one_hot_vector(3, ACTION_RIGHT_VECTOR_INDEX)
ACTION_UP_DOWN_NONE_VECTOR = make_one_hot_vector(3, ACTION_UP_DOWN_NONE_VECTOR_INDEX)
ACTION_LEFT_RIGHT_NONE_VECTOR = make_one_hot_vector(3, ACTION_LEFT_RIGHT_NONE_VECTOR_INDEX)
ACTION_UP_DOWN_TO_VECTOR_INDEX = {
    ACTION_UP: ACTION_UP_VECTOR_INDEX,
    ACTION_DOWN: ACTION_DOWN_VECTOR_INDEX,
    ACTION_UP_DOWN_NONE: ACTION_UP_DOWN_NONE_VECTOR_INDEX
}
ACTION_LEFT_RIGHT_TO_VECTOR_INDEX = {
    ACTION_LEFT: ACTION_LEFT_VECTOR_INDEX,
    ACTION_RIGHT: ACTION_RIGHT_VECTOR_INDEX,
    ACTION_LEFT_RIGHT_NONE: ACTION_LEFT_RIGHT_NONE_VECTOR_INDEX
}
ACTION_UP_DOWN_FROM_VECTOR_INDEX = [
    ACTION_UP,
    ACTION_DOWN,
    ACTION_UP_DOWN_NONE
]
ACTION_LEFT_RIGHT_FROM_VECTOR_INDEX = [
    ACTION_LEFT,
    ACTION_RIGHT,
    ACTION_LEFT_RIGHT_NONE
]
ACTION_UP_DOWN_TO_VECTOR = {
    ACTION_UP: ACTION_UP_VECTOR,
    ACTION_DOWN: ACTION_DOWN_VECTOR,
    ACTION_UP_DOWN_NONE: ACTION_UP_DOWN_NONE_VECTOR
}
ACTION_LEFT_RIGHT_TO_VECTOR = {
    ACTION_LEFT: ACTION_LEFT_VECTOR,
    ACTION_RIGHT: ACTION_RIGHT_VECTOR,
    ACTION_LEFT_RIGHT_NONE: ACTION_LEFT_RIGHT_NONE_VECTOR
}
ALL_MULTIACTIONS = []
ACTIONS_TO_MULTIVEC = dict()
ACTIONS_FROM_MULTIVEC_INDEX = []
def make_one_hot_vector(size, idx_one):
def get_random_action_up_down():
def get_random_action_left_right():
def keys_to_action_up_down(keys):
def keys_to_action_left_right(keys):
def action_up_down_to_vector(action):
def action_left_right_to_vector(action):
def action_up_down_to_vector_index(action):
def action_left_right_to_vector_index(action):
def action_up_down_from_vector_index(index):
def action_left_right_from_vector_index(index):
def action_to_key(action):
def actions_to_multivec(action_up_down, action_left_right): class Config(object):
    """Class containing some global configuration constants."""

    # path to the directory of the project
    MAIN_DIR = os.path.dirname(__file__)

    # GPU id to use
    GPU = 0

    # height and width at which ingame screenshots are saved to memory
    MODEL_HEIGHT = 720 // 4 # 180
    MODEL_WIDTH = 1280 // 4 # 320

    # directory of the annotations
    ANNOTATIONS_DIR = "/media/aj/grab/ml/ets2ai/annotations" if os.path.isdir("/media/aj/grab/ml/ets2ai/annotations") else os.path.join(MAIN_DIR, "annotations/")
    # filepath to the speed image annotations
    SPEED_IMAGE_SEGMENTS_DB_FILEPATH = os.path.join(MAIN_DIR, "annotations/speed_image_segments.txt")

    # directory where the replay memory files are saved
    REPLAY_MEMORIES_DIR = MAIN_DIR

    # configurations of the replay memories
    REPLAY_MEMORY_CFGS = {
        "reinforced-train": {"filepath": os.path.join(REPLAY_MEMORIES_DIR, "replay_memory_reinforced.sqlite"), "max_size": 650*1000, "max_size_tolerance": 20000},
        "reinforced-val": {"filepath": os.path.join(REPLAY_MEMORIES_DIR, "replay_memory_reinforced_val.sqlite"), "max_size": 150*1000, "max_size_tolerance": 20000},
        "supervised-train": {"filepath": os.path.join(REPLAY_MEMORIES_DIR, "replay_memory_supervised.sqlite"), "max_size": 200*1000, "max_size_tolerance": 20000},
        "supervised-val": {"filepath": os.path.join(REPLAY_MEMORIES_DIR, "replay_memory_supervised_val.sqlite"), "max_size": 50*1000, "max_size_tolerance": 20000},
    }

    # max speed, model inputs are clipped to that maximum
    MAX_SPEED = 150

    # max/min steering wheel positions to use for the classical and CNN-based
    # steering wheel prediction. values beyond this are viewed as somehow wrong
    # and either ignored or lead to the assumption that the current approximated
    # position is wrong.
    STEERING_WHEEL_MAX = 360+90+40 # wheel can be turned by roughly +/- 360+90 degrees, add 40deg tolerance
    STEERING_WHEEL_MIN = -STEERING_WHEEL_MAX # steering_wheel.py currently only looks at max value
    STEERING_WHEEL_RAW_ONE_MAX = 180
    STEERING_WHEEL_RAW_ONE_MIN = -STEERING_WHEEL_RAW_ONE_MAX
    STEERING_WHEEL_RAW_TWO_MAX = STEERING_WHEEL_RAW_ONE_MAX
    STEERING_WHEEL_RAW_TWO_MIN = STEERING_WHEEL_RAW_ONE_MIN
    STEERING_WHEEL_CNN_MAX = 360+90+40
    STEERING_WHEEL_CNN_MIN = -STEERING_WHEEL_CNN_MAX
    STEERING_WHEEL_RAW_CNN_MAX = 180
    STEERING_WHEEL_RAW_CNN_MIN = -STEERING_WHEEL_RAW_CNN_MAX

    # min and maximum direct reward range, used for computing bins
    MIN_REWARD = -100
    MAX_REWARD = 100

    # keys for pause and quickload to use
    KEY_PAUSE = "F1"
    KEY_QUICKLOAD = "F11"

    # Number of savegames to use during training. If set to N, the AI will
    # load a random one of the first N savegames during reinforcement learning.
    # (This happens many times during the training.)
    RELOAD_MAX_SAVEGAME_NUMBER = 6

    # Directory of an example image from an offence message.
    OFFENCE_FF_IMAGE = os.path.join(MAIN_DIR, "images/offence_ff.png")
    # Directory of an example image from a damage message (XX% damage).
    DAMAGE_SINGLE_DIGIT_IMAGE = os.path.join(MAIN_DIR, "images/damage_single_digit.png")
    # Directory of an example image from a damage message (X% damage).
    DAMAGE_DOUBLE_DIGIT_IMAGE = os.path.join(MAIN_DIR, "images/damage_double_digit.png")
    # Directory of an example image showing and activated reverse gear.
    REVERSE_IMAGE = os.path.join(MAIN_DIR, "images/reverse.png")
import sys
import os
import imgaug as ia
import cPickle as pickle
import argparse
import numpy as np
import cv2
import gzip as gz
import torch
import torch.nn.functional as F
import collections
import glob
        import scipy.ndimage as ndimage
        import scipy.ndimage.filters as filters
from train_semisupervised.train import (
    PREVIOUS_STATES_DISTANCES,
    MODEL_HEIGHT,
    MODEL_WIDTH,
    MODEL_PREV_HEIGHT,
    MODEL_PREV_WIDTH
)
from batching import states_to_batch
from train_semisupervised import models as models_semisupervised
from lib.util import to_variable, to_cuda, to_numpy
from lib import util
from lib import actions as actionslib
from config import Config
from imgaug import augmenters as iaa
from scipy import misc, ndimage
from skimage import draw
from __future__ import division, print_function

sys.path.append(os.path.join(os.path.dirname(__file__), '..'))



try:
    xrange
except NameError:
    xrange = range

def main():
    parser = argparse.ArgumentParser(description="Generate frames for a video. This requires a previous recording done via train.py --record='filename'.")
    parser.add_argument("--record", default=None, help="Filepath to recording.", required=True)
    parser.add_argument("--outdir", default="video_output", help="Output directory name.", required=False)
    parser.add_argument("--y", default=False, action="store_true", help="Always let ffmpeg/avconv overwrite existing video files without asking.", required=False)
    parser.add_argument("--noredraw", default=False, action="store_true", help="Skip drawing frames for videos for which the target directory already exists.", required=False)
    args = parser.parse_args()
    assert args.record is not None

    if "*" in args.record:
        record_fps = glob.glob(args.record)
        assert len(record_fps) > 0
        for record_fp in record_fps:
            fn = os.path.basename(record_fp)
            assert fn != ""
            fn_root = fn[0:fn.index(".")] #os.path.splitext(fn)
            assert fn_root != ""
            outdir = os.path.join(args.outdir, fn_root + "-" + str(abs(hash(record_fp)))[0:8])

            process(record_fp, outdir, args.y, args.noredraw)
    else:
        process(args.record, args.outdir, args.y, args.noredraw)

def process(record_fp, outdir_frames, y, noredraw):
    if outdir_frames[-1] == "/":
        outdir_videos = os.path.dirname(outdir_frames[:-1])
        dirname_frames_last = os.path.basename(outdir_frames[:-1])
    else:
        outdir_videos = os.path.dirname(outdir_frames)
        dirname_frames_last = os.path.basename(outdir_frames)
    assert os.path.isfile(record_fp)

    print("Processing recording '%s'..." % (record_fp,))
    print("Writing frames to '%s', videos to '%s' with filename start '%s'..." % (outdir_frames, outdir_videos, dirname_frames_last))

    with gz.open(record_fp, "rb") as f:
        recording = pickle.load(f)

    if os.path.exists(outdir_frames) and noredraw:
        print("Video frames were already drawn, not redrawing")
    else:
        if not os.path.exists(outdir_frames):
            print("Target directory for frames does not exist, creating it...")
            os.makedirs(outdir_frames)

        fd = FrameDrawer(outdir_videos)
        for fidx, frames_drawn in enumerate(fd.draw_frames(recording)):
            print("Frame %06d of around %06d..." % (fidx, len(recording["frames"])))
            frame_plans = frames_drawn[0]
            frame_atts = frames_drawn[1]
            frame_grids = frames_drawn[2]
            if frame_plans is not None:
                misc.imsave(os.path.join(outdir_frames, "plans_%06d.jpg" % (fidx,)), frame_plans)
            if frame_atts is not None:
                misc.imsave(os.path.join(outdir_frames, "atts_%06d.jpg" % (fidx,)), frame_atts)
            if frame_grids is not None:
                misc.imsave(os.path.join(outdir_frames, "grids_%06d.jpg" % (fidx,)), frame_grids)
            #if fidx > 200:
            #    break

    if not os.path.exists(outdir_videos):
        print("Target directory for videos does not exist, creating it...")
        os.makedirs(outdir_videos)

    frame_fps = ["plans_%06d.jpg", "atts_%06d.jpg", "grids_%06d.jpg"]
    frame_fps = [os.path.join(outdir_frames, fp) for fp in frame_fps]
    video_fps = ["plans.mp4", "atts.mp4", "grids.mp4"]
    video_fps = [os.path.join(outdir_videos, "%s-%s" % (dirname_frames_last, fp)) for fp in video_fps]

    for frame_fp, video_fp in zip(frame_fps, video_fps):
        #os.system('avconv %s -framerate 10 -i "%s" -crf 25 -b:v 2000k -vcodec mpeg4 %s' % ("-y" if y else "", frame_fp, video_fp))
        os.system('avconv %s -framerate 10 -i "%s" -crf 25 -b:v 2000k -vcodec h264 %s' % ("-y" if y else "", frame_fp, video_fp))

class FrameDrawer(object):
    def __init__(self, outdir):
        self.outdir = outdir

        checkpoint_supervised = torch.load("../train_semisupervised/train_semisupervised_model_withshortcuts.tar")
        embedder_supervised = models_semisupervised.PredictorWithShortcuts()
        embedder_supervised.eval()
        embedder_supervised.load_state_dict(checkpoint_supervised["predictor_state_dict"])

        if Config.GPU >= 0: