# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import os
os.environ["CUDA_VISIBLE_DEVICES"]="0"
os.environ["OMP_NUM_THREADS"]="1"
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import numpy as np
import math
import sys
import glob
from tensorflow.keras.utils import to_categorical
import time
import pickle
exp_name = 'track_seq4_vin_intsum_predv'

# # DATA

# +
v_max = 0
v_min = 0
def read_data_n_label(data, label):
    global v_max, v_min
    input_data = np.load(data)[:,[0,1,3,2,9]]
    input_data[:,2] = np.log(input_data[:,2])/16.
    input_data = input_data[input_data[:,0] > 0]
    input_data = input_data[input_data[:,0] < 60]
    input_data = input_data[input_data[:,1] > -40]
    input_data = input_data[input_data[:,1] < 40]

    grid_in = np.zeros((1,600,800,2))
    count_in = np.ones((1,600,800,1))

    for inz in input_data:
        pos_x = int(inz[0]//0.1)
        pos_y = int((inz[1] + 40)//0.1)
        grid_in[0, pos_x, pos_y, 0] += inz[2]#(grid_in[0, pos_x, pos_y, 0] + inz[2])#/(count_in[0, pos_x, pos_y, 0])
        grid_in[0, pos_x, pos_y, 1] = (grid_in[0, pos_x, pos_y, 1]*(count_in[0, pos_x, pos_y, 0]-1) + inz[3])/count_in[0, pos_x, pos_y, 0]
        count_in[0, pos_x, pos_y, 0] += 1

    with open(label) as f:
        content = f.readlines()
    content = [x.strip() for x in content]
    grid_label = np.zeros((1,300,400,8))
    if len(content) == 0:
        return grid_in, grid_label, np.array([-1,-1,-1,-1])
    velocity = (input_data[:,3][input_data[:,-1]>=0])
    velocity = velocity.mean() if velocity.size>0 else np.zeros(1)
    velocity = velocity/4
    tar_location = np.fromstring(content[1][1:-1], dtype= float, sep = ' ')
    tar_WL = np.fromstring(content[2][1:-1], dtype= float, sep = ',')
    tar_angle = float(content[3])

    pos0 = int((tar_location[0])//0.2)
    pos1 = int((tar_location[1] + 40)//0.2)

    grid_label[0, pos0, pos1, -1] = 1
    grid_label[0, pos0, pos1, -2] = tar_angle
    grid_label[0, pos0, pos1, 3] = tar_WL[0]
    grid_label[0, pos0, pos1, 4] = tar_WL[1]
 
    
    if tar_WL[0] < tar_WL[1]:
        grid_label[0, pos0, pos1, 3] = tar_WL[1]
        grid_label[0, pos0, pos1, 4] = tar_WL[0]
        grid_label[0, pos0, pos1, -2] = tar_angle + np.pi/2
#         print ('ALARM')
        
    grid_label[0, pos0, pos1, 1] = tar_location[0] - (pos0*0.2) - 0.1
    grid_label[0, pos0, pos1, 2] = tar_location[1] - (pos1*0.2 -40) -0.1
    grid_label[0, pos0, pos1, 5] = velocity
    if content[0] == 'bike':
        grid_label[0, pos0, pos1, 0]  = 1
    elif content[0] == 'pedestrian':
        grid_label[0, pos0, pos1, 0]  = 2
    else:
        grid_label[0, pos0, pos1, 0]  = 0
    if grid_label[0, pos0, pos1, -2]>(v_max):
#         v_max=grid_label[0, pos0, pos1, -2]
        v_max=grid_in[0, :, :, 0].max()
    if grid_label[0, pos0, pos1, -2]<(v_min):
#         v_min=grid_label[0, pos0, pos1, -2]
        v_min=grid_in[0, :, :, 0].min()
    return grid_in, grid_label, np.array([tar_location[0], tar_location[1], tar_WL[0], tar_WL[1]])


# X.shape, y[:,:,:,0].shape
# -

training_folders = glob.glob('training_data/*')
training_labels = glob.glob(training_folders[0] + '/labels/radar_left_np/*')
training_X = glob.glob(training_folders[0] + '/radar_left_np/*')
training_labels.sort()
training_X.sort()
X = []
y = []
for j in range(len(training_folders)):
    training_labels = glob.glob(training_folders[j] + '/labels/radar_left_np/*')
    training_X = glob.glob(training_folders[j] + '/radar_left_np/*')
#     print (j,training_folders[j])
    for i in range(len(training_labels)):
        X_in, y_in, location = read_data_n_label(training_X[i], training_labels[i])
        X.append(X_in)
        y.append(y_in)
X = np.array(X)
X = X.reshape(len(X),600,800,2)
y = np.array(y)
y = y.reshape(len(y),300,400,8)


def feature_extractor(image_in):
    
    with tf.variable_scope('feature_extractor'):
        net = tf.layers.conv2d(image_in, 8, 3,strides=2, activation=tf.nn.relu, padding='same')
        net = tf.layers.conv2d(net, 8, 3, strides=1,activation=tf.nn.relu, padding='same')
        net = tf.layers.conv2d(net, 8, 3, strides=1,activation=tf.nn.relu, padding='same')

        x1 = net
        net = tf.layers.conv2d(net, 16, 3,strides=2, activation=tf.nn.relu, padding='same')
        net = tf.layers.conv2d(net, 32, 3, strides=1,activation=tf.nn.relu, padding='same')
        net = tf.layers.conv2d(net, 32, 3, strides=1,activation=tf.nn.relu, padding='same')

        x2 = net
    return x1, x2


# +
images_in = tf.placeholder(tf.float32, (None,600,800,2,4))
y_label = tf.placeholder(tf.int32, (None,300, 400, 4, 1 ))
y_label_one_hot = tf.one_hot(y_label, 3, axis=-1)
y_angle = tf.placeholder(tf.float32, (None,300, 400 , 4, 1))
y_wl = tf.placeholder(tf.float32, (None,300, 400 , 4, 2))
y_loc = tf.placeholder(tf.float32, (None,300, 400 , 4, 2))
y_v = tf.placeholder(tf.float32, (None,300, 400 , 4, 1))

y_mask = tf.placeholder(tf.float32, (None,300, 400 , 4))

with tf.variable_scope("features") as scope:
    x11, x12 = feature_extractor(images_in[:,:,:,:,0])
    scope.reuse_variables()
    x21, x22 = feature_extractor(images_in[:,:,:,:,1])
    scope.reuse_variables()
    x31, x32 = feature_extractor(images_in[:,:,:,:,2])
    scope.reuse_variables()
    x41, x42 = feature_extractor(images_in[:,:,:,:,3])
    
x1 = tf.concat([x11, x21, x31, x41], axis = -1)
x2 = tf.concat([x12, x22, x32, x42], axis = -1)
up1 = tf.layers.conv2d_transpose(x1, 32, 3, activation=tf.nn.relu, padding='same')
up2 = tf.layers.conv2d_transpose(x2, 64, 3,strides=2, activation=tf.nn.relu, padding='same')

concat = tf.concat([up1, up2], axis = -1)

# -

def prediction_head(concat_vec, y_label_one_hot_ts, y_angle_ts, y_loc_ts, y_wl_ts, y_v_ts, y_mask_ts, name):
    
    with tf.variable_scope(name):
        y_pred_occ = tf.layers.conv2d(concat_vec, 3 , 1, activation=None, padding= 'same')
        y_pred_occ_prob = tf.nn.softmax(y_pred_occ, axis = -1)
        y_pred_angle = tf.layers.conv2d(concat_vec, 1 , 1, activation=None, padding= 'same')
        y_pred_loc   = tf.layers.conv2d(concat_vec, 2 , 1, activation=None, padding= 'same')
        y_pred_WL   = tf.layers.conv2d(concat_vec, 2 , 1, activation=tf.nn.softplus, padding= 'same')
        y_pred_v = tf.layers.conv2d(concat_vec, 1 , 1, activation=None, padding= 'same')
        
        
        log_y_pred_WL = tf.log(y_pred_WL + 0.1)
        log_y_wl = tf.log(y_wl_ts + 0.1)

        loss_occ = tf.reduce_mean(-y_label_one_hot_ts[:,:, : , 0] *0.008*((1 - y_pred_occ_prob[:,:,:,0])) * tf.log(y_pred_occ_prob[:,:,:,0]  + 0.01 )\
                             -y_label_one_hot_ts[:,:, : , 1] *0.992*((1 - y_pred_occ_prob[:,:,:,1])) * tf.log(y_pred_occ_prob[:,:,:,1] + 0.01)\
                              -y_label_one_hot_ts[:,:, : , 2] *0.992*((1 - y_pred_occ_prob[:,:,:,2])) * tf.log(y_pred_occ_prob[:,:,:,2] + 0.01)\
                             )
        loss_angle = tf.reduce_mean(y_mask_ts*tf.abs(tf.sin(y_pred_angle - y_angle_ts)))
        loss_loc = tf.reduce_mean(y_mask_ts*tf.abs(y_loc_ts- y_pred_loc))
        loss_wl = tf.reduce_mean(y_mask_ts*tf.abs(log_y_wl - log_y_pred_WL))
        loss_v = tf.reduce_mean(y_mask_ts*tf.abs(y_v_ts - y_pred_v)*(y_v_ts!=0))


        total_loss = 1.0*loss_occ + 1.25*loss_angle + 0.75*loss_loc + loss_wl + 1 * loss_v
        
    return y_pred_occ_prob, y_pred_angle, y_pred_loc, y_pred_WL, y_pred_v, total_loss

# +
y_pred_occ_prob0, y_pred_angle0, y_pred_loc0, y_pred_WL0, y_pred_v0, total_loss0 = prediction_head(concat, y_label_one_hot[:,:,:,0,0,:], y_angle[:,:,:,0,:], \
                y_loc[:,:,:,0,:], y_wl[:,:,:,0,:], y_v[:,:,:,0,:], tf.expand_dims(y_mask[:,:,:,0], axis = -1), 'prediction_0')

y_pred_occ_prob1, y_pred_angle1, y_pred_loc1, y_pred_WL1, y_pred_v1, total_loss1 = prediction_head(concat, y_label_one_hot[:,:,:,1,0,:], y_angle[:,:,:,1,:], \
                y_loc[:,:,:,1,:], y_wl[:,:,:,1,:], y_v[:,:,:,1,:], tf.expand_dims(y_mask[:,:,:,1], axis = -1), 'prediction_1')

y_pred_occ_prob2, y_pred_angle2, y_pred_loc2, y_pred_WL2, y_pred_v2, total_loss2 = prediction_head(concat, y_label_one_hot[:,:,:,2,0,:], y_angle[:,:,:,2,:], \
                y_loc[:,:,:,2,:], y_wl[:,:,:,2,:], y_v[:,:,:,2,:], tf.expand_dims(y_mask[:,:,:,2], axis = -1), 'prediction_2')

y_pred_occ_prob3, y_pred_angle3, y_pred_loc3, y_pred_WL3, y_pred_v3, total_loss3 = prediction_head(concat, y_label_one_hot[:,:,:,3,0,:], y_angle[:,:,:,3,:], \
                y_loc[:,:,:,3,:], y_wl[:,:,:,3,:], y_v[:,:,:,3,:], tf.expand_dims(y_mask[:,:,:,3], axis = -1), 'prediction_3')

total_loss = total_loss0 + 0.75*total_loss1 + 0.5*total_loss2 + 0.25*total_loss3


# -

v_min,v_max #angle(-1.5687682499618538, 1.5707963267948966)
#velocity (-4.0950820584830465, 4.674418471324403)
#intensity (0.0, 4.248383107944953)

# +
def processing_seq_data(X, Y):
    b_x = []
    b_y = []
    Y = np.expand_dims(Y, axis = 3)
    X = np.expand_dims(X, axis = -1)
    for i in range(len(X) -6):
        b_x.append(np.concatenate((X[i], X[i+1], X[i+2], X[i+3]), axis = -1))
        b_y.append(np.concatenate((Y[i+3], Y[i+4], Y[i+5], Y[i+6]), axis = 2))
    
    b_x = np.asarray(b_x)
    b_y = np.asarray(b_y)
    
    return b_x, b_y

def dataset(np_seed, seqlen=10):
    Xs = []
    labels = []
    training_folders = glob.glob('training_data/*')
    for j in range(len(training_folders)):
        training_labels = glob.glob(training_folders[j] + '/labels/radar_left_np/*')
        training_X = glob.glob(training_folders[j] + '/radar_left_np/*')
        training_labels.sort()
        training_X.sort()
        i_begin = np.random.RandomState(np_seed).randint(seqlen)
        training_labels = training_labels[i_begin:]
        training_X = training_X[i_begin:]
        n_seq = len(training_X)//seqlen
        i_end = n_seq * seqlen
        training_labels = training_labels[:i_end]
        training_X = training_X[:i_end]
#         print(len(training_labels))
        Xs.extend(training_X)
        labels.extend(training_labels)
    return Xs, labels


# +
#shuffle 1
np_seed = 12345
optimizer = tf.train.AdamOptimizer().minimize(total_loss)
init = tf.global_variables_initializer()
saver = tf.train.Saver(max_to_keep=20)
seqlen = 10

with tf.Session() as sess:
    sess.run(init)
    for i_epoch in range(30):
        batch_X = []
        batch_y = []
        t_begin = time.time()
        Xs, labels = dataset(np_seed+i_epoch,seqlen=seqlen)
        n_seq = len(Xs)
        print(n_seq)
        seq_shuffle = np.random.RandomState(np_seed+i_epoch).permutation(n_seq//seqlen)
        for i, i_seq in enumerate(seq_shuffle):
            for ii in range(seqlen):
                X_in, y_in, _ = read_data_n_label(Xs[i_seq*seqlen+ii], labels[i_seq*seqlen+ii])
#                 print (i,' ',i_seq,' ',ii,' ',i_seq*seqlen+ii)
#                 print(Xs[i_seq*seqlen+ii])
                batch_X.append(X_in)
                batch_y.append(y_in)

            batch_X_seq = np.array(batch_X).reshape(seqlen,600,800,2)
            batch_y_seq = np.array(batch_y).reshape(seqlen,300,400,8)

            batch_X, batch_y = processing_seq_data(batch_X_seq, batch_y_seq)

            pred_occ, cost,_= sess.run([y_pred_occ_prob0,total_loss, optimizer], feed_dict = {images_in: batch_X,\
                                                        y_label: batch_y[:,:,:,:,0:1], \
                                                        y_angle: batch_y[:,:,:,:,-2:-1],\
                                                        y_loc: batch_y[:,:,:,:,[1,2]], \
                                                        y_wl: batch_y[:,:,:,:,[3,4]], \
                                                        y_v: batch_y[:,:,:,:,5:6],\
                                                        y_mask: batch_y[:,:,:,:,-1]})

            if i_seq % (n_seq//4)==0:
                print('epoch {}'.format(i_epoch))
                print ('cost: ', cost)
                z = pred_occ[-1]
                argmax = z.argmax(axis = 2)
                print ('--------------------------------------------------------')
                print ('Predictions:',np.where(argmax > 0))
                print ('labels:', np.where(batch_y[-1,:,:,0,-1] > 0.5))
                print ('Pre_ID: ', np.argmax(pred_occ[-1]),', GT_ID: ', np.argmax(batch_y[-1,:,:,-1]))
            batch_X = []
            batch_y = []
        save_path = saver.save(sess, "saved_models/{}/model_e{}.ckpt".format(exp_name, i_epoch))
        t_elapsed = time.time()-t_begin
        print('Time Elapsed: ', t_elapsed)
    save_path = saver.save(sess, "saved_models/{}/model.ckpt".format(exp_name))
# -

# ## Save Testing Predictions 

# exp_name='track-predv0'
for i_epoch in range(11,30):
    saver = tf.train.Saver()
    t_begin = time.time()
    with tf.Session() as sess:
        saver.restore(sess, "saved_models/{}/model_e{}.ckpt".format(exp_name, i_epoch))
    #     saver.restore(sess, "saved_models/model_e{}.ckpt".format( i_epoch))
        print("Model restored.")
        predictions = {'metadata': [],'label_preds':[], 'scores':[], 'box3d_lidar':[],'velocity':[],'velocity_gt':[]}
        counter = 0
        training_folders = glob.glob('validation_data/*')
        training_folders.sort()
        for folder in range(len(training_folders)):
            training_labels = glob.glob(training_folders[folder] + '/labels/radar_left_np/*')
            print(folder)
            training_X = glob.glob(training_folders[folder] + '/radar_left_np/*')
            training_labels.sort()
            training_X.sort()
            X = []
            y = []
            for i in range(len(training_labels)):
                X_in, y_in, location = read_data_n_label(training_X[i], training_labels[i])
                X.append(X_in)
                y.append(y_in)
            X = np.concatenate(X, axis=0)
            y = np.concatenate(y, axis=0)
            X = np.concatenate([np.tile(X[:1],(6,1,1,1)),X],axis=0)
            var_names = ['occ','angle','loc','wl', 'v']
            for j in range(4):
                for var in var_names:
                    code = "pred_{}X{}s=[]".format(var,j)
                    exec(code)       
            for i in range(3,len(X)):
                pred_occX0, pred_angleX0, pred_locX0, pred_wlX0,pred_vX0, pred_occX1, pred_angleX1, pred_locX1, pred_wlX1,pred_vX1,\
                pred_occX2, pred_angleX2, pred_locX2, pred_wlX2,pred_vX2, pred_occX3, pred_angleX3, pred_locX3, pred_wlX3,pred_vX3 \
                                                                = sess.run([y_pred_occ_prob0, y_pred_angle0, y_pred_loc0, y_pred_WL0,y_pred_v0,\
                                                                           y_pred_occ_prob1, y_pred_angle1, y_pred_loc1, y_pred_WL1,y_pred_v1,\
                                                                           y_pred_occ_prob2, y_pred_angle2, y_pred_loc2, y_pred_WL2,y_pred_v2,\
                                                                           y_pred_occ_prob3, y_pred_angle3, y_pred_loc3, y_pred_WL3,y_pred_v3], \
                                                               feed_dict = {images_in: np.stack((X[i-3],X[i-2],X[i-1],X[i]),axis=-1)[np.newaxis]})
                for j in range(4):
                    for var in var_names:
                        code = "pred_{}X{}s.append(pred_{}X{})".format(var,j,var,j)
                        exec(code)
            for i in range(6, len(X)):
                for var in var_names:
                    for j in range(4):
                        code = "pred_{}X{}=pred_{}X{}s[{}]".format(var,j,var,j,i-3-j)
                        exec(code)

                    code = "pred_{} = np.mean(np.array([pred_{}X0, pred_{}X1, pred_{}X2, pred_{}X3]), axis = 0)".format(var,var,var,var,var)
                    exec(code)

                pred_index = np.asarray(np.where(pred_occ[0,:,:,:].argmax(axis = -1)>0)).T
                #pred_index = np.asarray(np.where(pred_occ[0,:,:,0] < 0.5)).T
                iii = i -6
                try:
                    v_gt = y[iii][...,5].sum()
                except:
                    v_gt = 0
                predictions['velocity_gt'].append(v_gt*4)
                predictions['metadata'].append(training_X[iii])

                predictions['label_preds'].append([])
                predictions['box3d_lidar'].append([])
                predictions['scores'].append([])
                predictions['velocity'].append([])

                for index in range(len(pred_index)):

                    predictions['label_preds'][counter].append(pred_occ[:,pred_index[index,0], pred_index[index,1], [1,2]].argmax(axis = -1)[0])
                    predictions['scores'][counter].append(pred_occ[:,pred_index[index,0], pred_index[index,1], [1,2]].max(axis = -1)[0])
                    predictions['velocity'][counter].append(pred_v[:,pred_index[index,0], pred_index[index,1], 0]*4)

                    pred_location_x =  pred_loc[:,pred_index[index,0], pred_index[index,1],0] + (pred_index[index,0]*0.2) + 0.1
                    pred_location_y =  pred_loc[:,pred_index[index,0], pred_index[index,1],1] + (pred_index[index,1]*0.2 -40) + 0.1

                    W_location_xy = pred_wl[:,pred_index[index,0], pred_index[index,1],0]
                    L_location_xy = pred_wl[:,pred_index[index,0], pred_index[index,1],1]

                    angle_xy = pred_angle[:,pred_index[index,0], pred_index[index,1],0]
                    predictions['box3d_lidar'][counter].append([])
                    predictions['box3d_lidar'][counter][-1].append(pred_location_x[0])
                    predictions['box3d_lidar'][counter][-1].append(pred_location_y[0])
                    predictions['box3d_lidar'][counter][-1].append(-1)
                    predictions['box3d_lidar'][counter][-1].append(W_location_xy[0])
                    predictions['box3d_lidar'][counter][-1].append(L_location_xy[0])
                    predictions['box3d_lidar'][counter][-1].append(-1)
                    predictions['box3d_lidar'][counter][-1].append(angle_xy[0])
                counter += 1
    fname = "saved_models/{}/model_e{}.pkl".format(exp_name, i_epoch)
    with open(fname, 'wb') as output:
        pickle.dump(predictions, output)
    print(time.time()-t_begin)

# +
# save NMS prediction
from shapely.geometry import Polygon
def get_rect(filler):
    W = filler[3]
    L = filler[4]
    angle = filler[-1]

    m0 = np.array([filler[0], filler[1]])
    k1 = np.array([np.cos(angle), np.sin(angle)])*W
    k2 = np.array([np.sin(angle), -np.cos(angle)])*L
    p0 = m0 - k1 - k2 # m0 + k1 + k2
    p1 = m0 +k1 - k2 #m0 - k1 - k2
    p2 = m0 + k1 + k2
    p3 = m0 - k1 + k2
    
    return Polygon([p0, p1, p2, p3]), [p0, p1, p2, p3]

# exp_name = 'track_seq4_vin_intsum_predv'
nonobject= True #False
for i_epoch in range(12,30):
    fname = "saved_models/{}/model_e{}".format(exp_name, i_epoch)
    with open(fname+".pkl", 'rb') as file:
        object_file = pickle.load(file)
    nms_dict = {'metadata': [],'label_preds':[], 'scores':[], 'box3d_lidar':[], 'velocity':[]}
    v_loss = []
    for i in range(len(object_file['metadata'])):

        score = np.array(object_file['scores'][i])
        llabel = np.array(object_file['label_preds'][i])
        rrect  = np.array(object_file['box3d_lidar'][i])
        velo = np.array(object_file['velocity'][i])
        """okay = np.where(score> 0.45)[0]
        score = score[okay]
        llabel = llabel[okay]
        rrect = rrect[okay]"""

        nms_dict['metadata'].append(object_file['metadata'][i])
        nms_dict['label_preds'].append([])
        nms_dict['box3d_lidar'].append([])
        nms_dict['scores'].append([])
        nms_dict['velocity'].append([])
        if len(score)==0:
            v_diff = object_file['velocity_gt'][i]
        else:
            v_diff = None
        while(len(score) > 0):
            nms_dict['box3d_lidar'][i].append([])
            under_attention = score.argmax()
            max_score = score.max()
    #         print(max_score)
            class_pred = llabel[under_attention]
            #score = np.delete(score, under_attention)
            if v_diff is None:
                if object_file['velocity_gt'][i]==0:
                    v_diff = 1
                else:
                    v_diff = (velo[under_attention] - object_file['velocity_gt'][i])/(object_file['velocity_gt'][i])
    #         print (llabel)
    #         print (rrect.shape)
    #         print(score.shape)
            box_attent, box_arr = get_rect(rrect[under_attention])
            box_apc = rrect[under_attention]
            #print (object_file['box3d_lidar'][i])
            #nms_dict['box3d_lidar'][i][-1].append(box_arr)
            box_arr = np.array(box_arr)
            deletion_list = []
            avg_conf= []
            count = 1.0
            for mmm in range(len(score)):
                boxmm, boxyy = get_rect(rrect[mmm])
                boxyy = np.array(boxyy)
                if (box_attent.intersection(boxmm).area / box_attent.union(boxmm).area) > 0:
                    if nonobject or (not nonobject and class_pred == llabel[mmm]):
                        deletion_list.append(mmm)

            nms_dict['label_preds'][i].append(class_pred)
            nms_dict['scores'][i].append(max_score)
            nms_dict['box3d_lidar'][i][-1].append(box_apc)
            nms_dict['velocity'][i].append(velo[under_attention])
            score = np.delete(score, deletion_list)
            rrect = np.delete(rrect, deletion_list, axis = 0)
            llabel = np.delete(llabel, deletion_list)
        v_loss.append(v_diff)
    v_loss = np.array(v_loss)
    v_loss1 = np.abs(v_loss).mean() #/len(object_file['metadata'])
    v_loss2 = (v_loss**2).mean() #/len(object_file['metadata'])
    print(i_epoch, 'velocity loss: ', v_loss1,' ', v_loss2)
    
    nms_list = []
    for i in range(len(nms_dict['metadata'])):
        sub_dict = dict()
        sub_dict['metadata'] = dict()
        sub_dict['metadata']['image_idx']= nms_dict['metadata'][i]
    #     sub_dict['metadata']['image_shape']= np.array([ 600, 800])
        sub_dict['label_preds'] = np.array(nms_dict['label_preds'][i])
        sub_dict['scores'] = np.array(nms_dict['scores'][i])
        sub_dict['velocity'] = np.array(nms_dict['velocity'][i])
        if len(nms_dict['box3d_lidar'][i]) >0:
            sub_dict['box3d_lidar'] = np.array(nms_dict['box3d_lidar'][i]).squeeze(axis = 1)
        else:
            sub_dict['box3d_lidar'] = np.array(nms_dict['box3d_lidar'][i])
        nms_list.append(sub_dict)
    
    if nonobject:
        fname +='-n'
    with open(fname+'-nms.pkl', 'wb') as output:
        pickle.dump(nms_list, output)

StarcoderdataJupyterScriptsDedupFiltered