visualization_side.py

import os
import pickle
import numpy as np
import pandas as pd
import shutil
from src.utils.visualization import create_gif, draw_belief, plot_reward, propagate_belief
from src.intersection.reward import IndependentReward

# Load Meta data
Meta_file = "Results_intersection" + os.sep + "Setups_intersection.xlsx"
Meta = pd.read_excel(Meta_file, index_col=0, keep_default_na=False)

# Go throught the results directory
result_files = os.listdir("Results_intersection")

# Exclude .xlsx files
result_files = [file for file in result_files if not file.endswith('.xlsx')]
result_files = ['Exp_0', 'Exp_1', 'Exp_2', 'Exp_3', 'Exp_4', 'Exp_5'] + result_files

create_gifs = True
add_preds_to_gif = True

ignore_unchanged = True
for result_file in result_files:
    # Check if the gif file exist, and if they are newer than the .pkl file,'
    # if so, skip this iteration
    if ".xlsx" in result_file:
        continue

    # Exclude pdf files
    if ".pdf" in result_file:
        continue

    # Exclude svg files
    if ".svg" in result_file:
        continue

    # Exclude npy files
    if ".npy" in result_file:
        continue

    # Exclude odt files
    if ".odt" in result_file:
        continue

    # Exclude odt files
    if ".csv" in result_file:
        continue

    # Exclude odt files
    if ".png" in result_file:
        continue

    # Get experiment number
    exp_index = int(result_file.split('_')[-1])
    meta = Meta.loc[exp_index]

    folder_path = "Results_intersection" + os.sep + result_file + os.sep
    folder_files = os.listdir(folder_path)

    # Define reward function
    reward = IndependentReward(lane_width = meta["lane_width"], 
                               weigh_particles=meta["weigh_particles"],
                               full_violation_factor=meta["full_violation_factor"],
                               d = meta["d"])
    
    print('Update visualization for {}'.format(result_file))

    # Load the results
    with open(folder_path + result_file + '.pkl', 'rb') as f:
        data = pickle.load(f)

    # Get driven trajectories
    Eta = data['eta']
    A_cont = data['a_cont']
    A_cont_init = data['a_cont_init']
    # Check if multiple timesteps were saved
    if len(A_cont.shape) == 4:
        a_cont = A_cont[0]
    else:
        a_cont = A_cont
        A_cont = A_cont[np.newaxis]
    Eta = np.concatenate((Eta[...,:10], a_cont.transpose(1,0,2), Eta[...,10:]), axis = -1)

    Traj_ego = Eta[...,[0,1,2,3,4,10,11]] # [x, y, theta, delta, v, a, w]
    Traj_tar = Eta[...,[5,6,7,8,9,12,13]] # [x, y, theta, delta, v, a, w]

    B = data['b']
    S = np.concatenate((B[...,:12], np.ones((*B.shape[:3], 2)) * np.nan, B[...,12:]), axis = -1)
    W = data['w'][...,np.newaxis]
    Traj_ego_belief = S[...,[2,3,4,5,6,12,13]] # [x, y, theta, delta, v]
    Traj_tar_belief = S[...,[7,8,9,10,11,14,15]] # [x, y, theta, delta, v, a, w]

    num_cases, num_steps = Traj_ego.shape[:2]

    vehicle_length = 4.3
    vehicle_width = 1.72

    # Get radians to degrees
    Traj_ego[...,2] = Traj_ego[...,2] * 180 / np.pi
    Traj_tar[...,2] = Traj_tar[...,2] * 180 / np.pi
    Traj_ego_belief[...,2] = Traj_ego_belief[...,2] * 180 / np.pi
    Traj_tar_belief[...,2] = Traj_tar_belief[...,2] * 180 / np.pi

    # get rewards
    Reward = data['v']
    Reward_init = data['v_init']

    # Get the modification time of the pkl file
    pkl_mod_time = os.path.getmtime(folder_path + result_file + '.pkl')

    for case in range(num_cases):
        # Test time
        gif_files = [file for file in folder_files if (file.endswith('.mp4') and 'case=' + str(case) in file)]
        pdf_files = [file for file in folder_files if (file.endswith('.pdf') and 'case=' + str(case) in file)]
        
        if len(gif_files) > 0:
            # Check the modification time of the gif files
            gif_mod_time = np.array([os.path.getmtime(folder_path + file) for file in gif_files])
            gif_mod_time = np.min(gif_mod_time)
        else:
            gif_mod_time = 0

        if len(pdf_files) > 0:
            # Check the modification time of the gif files
            pdf_mod_time = np.array([os.path.getmtime(folder_path + file) for file in pdf_files])
            pdf_mod_time = np.min(pdf_mod_time)
        else:
            pdf_mod_time = 0
        
        # Get mp4 file
        result_path = folder_path + os.sep + result_file
        
        # Create gifs
        if create_gifs: # and (not gif_mod_time > pkl_mod_time):
            a_cont = A_cont[:,:,[case]]
            if len(A_cont.shape) < 4:
                a_cont = np.concatenate([a_cont, np.zeros((19, *a_cont.shape[1:]))], axis = -2)
            a_cont = a_cont.transpose(2,1,0,3)
            if add_preds_to_gif:
                S_pred = propagate_belief(B[[case]], a_cont, meta, reward)
                S_pred = S_pred[0]
            else:
                S_pred = None

            create_gif(Traj_ego, Traj_tar, Traj_ego_belief, Traj_tar_belief, W, Reward, Reward_init, a_cont, case, meta, result_path, S_pred)
        
        if not pdf_mod_time > pkl_mod_time:
            # Plot belief states
            draw_belief(Eta, S, W, case, ['x_ego', 'x_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['y_ego', 'y_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['v_ego', 'v_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['theta_ego', 'theta_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['delta_ego', 'delta_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['steering_ego', 'steering_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['a_ego', 'a_tar'], meta, result_path)
            draw_belief(Eta, S, W, case, ['w_ego' ,'w_tar'], meta, result_path)
            
        # Plot rewards
        try:
            plot_reward(Reward, Reward_init, case, meta, result_path)
        except:
            pass

    # delete frames
    Frames_folder = folder_path + 'Frames'
    if os.path.exists(Frames_folder):
        shutil.rmtree(Frames_folder)