from pie_intent import *
from pie_data import PIE


data_opts = {
        'fstride': 1,
        'sample_type': 'all',
        'height_rng': [0, float('inf')],
        'squarify_ratio': 0,
        'data_split_type': 'default',  # kfold, random, default
        'seq_type': 'intention',  # crossing , intention
        'min_track_size': 0,  # discard tracks that are shorter
        'max_size_observe': 15,  # number of observation frames
        'max_size_predict': 5,  # number of prediction frames
        'seq_overlap_rate': 0.3,  # how much consecutive sequences overlap
        'balance': True,  # balance the training and testing samples
        'crop_type': 'context',  # crop 2x size of bbox around the pedestrian
        'crop_mode': 'pad_resize',  # pad with 0s and resize to VGG input
        'encoder_input_type': [],
        'decoder_input_type': ['bbox'],
        'output_type': ['intention_binary']
    }

t = PIEIntent(num_hidden_units=128,
                  regularizer_val=0.001,
                  lstm_dropout=0.4,
                  lstm_recurrent_dropout=0.2,
                  convlstm_num_filters=64,
                  convlstm_kernel_size=2)
imdb = PIE(data_path=os.environ.copy()['PIE_PATH'])

beh_seq_set04 = imdb.generate_data_trajectory_sequence('set04',**data_opts) #读取所有的原始数据标签

data_type = {'encoder_input_type': data_opts['encoder_input_type'],
             'decoder_input_type': data_opts['decoder_input_type'],
             'output_type': data_opts['output_type']}

t._model_type = 'convlstm_encdec'
seq_length = data_opts['max_size_observe']  # 15

# train_d = t.get_train_val_data(beh_seq_set04, data_type, seq_length, data_opts['seq_overlap_rate'])

occ=beh_seq_set04['occlusion']
img_sequences=beh_seq_set04['image']
bbox_sequences=beh_seq_set04['bbox']
ped_ids=beh_seq_set04['ped_id']

print(len(occ))

for i in range(len(occ)):
    #找到符合要求的索引
    id=ped_ids[i]
    if id[0][0]=='4_6_1264':
        img_list=img_sequences[i]
        occ_list=occ[i]
        frame=[]
        for imp in img_list:
            f = imp.split('/')[-1].split('.')[0]
            frame.append(f)
        for j in range(len(occ_list)):
            print(frame[j]+' : ',occ_list[j])


# i=-1
# for seq, pid in zip(img_sequences, ped_ids):
#     i += 1
#     update_progress(i / len(img_sequences))
#     img_seq = []
#     for imp, b, p in zip(seq, bbox_sequences[i], pid):
#         set_id = imp.split('/')[-3]
#         vid_id = imp.split('/')[-2]
#         img_name = imp.split('/')[-1].split('.')[0]
#         img_save_folder = os.path.join(save_path, set_id, vid_id)