106 lines
4.7 KiB
Python
106 lines
4.7 KiB
Python
|
import pandas as pd
|
||
|
|
import os
|
||
|
|
import glob
|
||
|
|
import pickle
|
||
|
|
|
||
|
|
DATASET_LOCATION = "YOUR_PATH_HERE"
|
||
|
|
|
||
|
|
def reframe_annotation():
|
||
|
|
annotation_path = f'{DATASET_LOCATION}/retrieve_annotation/all/'
|
||
|
|
save_path = f'{DATASET_LOCATION}/reformat_annotation/'
|
||
|
|
if not os.path.exists(save_path):
|
||
|
|
os.makedirs(save_path)
|
||
|
|
tasks = glob.glob(annotation_path + '*.txt')
|
||
|
|
id_map = pd.read_csv('id_map.csv')
|
||
|
|
for task in tasks:
|
||
|
|
if not task.split('/')[-1].split('_')[2] == '1.txt':
|
||
|
|
continue
|
||
|
|
with open(task, 'r') as f:
|
||
|
|
lines = f.readlines()
|
||
|
|
task_id = int(task.split('/')[-1].split('_')[1]) + 1
|
||
|
|
clip = id_map.loc[id_map['ID'] == task_id].folder
|
||
|
|
print(task_id, len(clip))
|
||
|
|
if len(clip) == 0:
|
||
|
|
continue
|
||
|
|
with open(save_path + clip.item() + '.txt', 'w') as f:
|
||
|
|
for line in lines:
|
||
|
|
words = line.split()
|
||
|
|
f.write(words[0] + ',' + words[1] + ',' + words[2] + ',' + words[3] + ',' + words[4] + ',' + words[5] +
|
||
|
|
',' + words[6] + ',' + words[7] + ',' + words[8] + ',' + words[9] + ',' + ' '.join(words[10:]) + '\n')
|
||
|
|
f.close()
|
||
|
|
|
||
|
|
def get_grid_location(obj_frame):
|
||
|
|
x_min = obj_frame['x_min']#.item()
|
||
|
|
y_min = obj_frame['y_min']#.item()
|
||
|
|
x_max = obj_frame['x_max']#.item()
|
||
|
|
y_max = obj_frame['y_max']#.item()
|
||
|
|
gridLW = 1280 / 25.
|
||
|
|
gridLH = 720 / 15.
|
||
|
|
center_x, center_y = (x_min + x_max)/2, (y_min + y_max)/2
|
||
|
|
X, Y = int(center_x / gridLW), int(center_y / gridLH)
|
||
|
|
return X, Y
|
||
|
|
|
||
|
|
def regenerate_annotation():
|
||
|
|
annotation_path = f'{DATASET_LOCATION}/reformat_annotation/'
|
||
|
|
save_path=f'{DATASET_LOCATION}/regenerate_annotation/'
|
||
|
|
if not os.path.exists(save_path):
|
||
|
|
os.makedirs(save_path)
|
||
|
|
tasks = glob.glob(annotation_path + '*.txt')
|
||
|
|
for task in tasks:
|
||
|
|
print(task)
|
||
|
|
annt = pd.read_csv(task, sep=",", header=None)
|
||
|
|
annt.columns = ["obj_id", "x_min", "y_min", "x_max", "y_max", "frame", "lost", "occluded", "generated", "name", "label"]
|
||
|
|
obj_records = {}
|
||
|
|
for index, obj_frame in annt.iterrows():
|
||
|
|
if obj_frame['name'].startswith('P'):
|
||
|
|
continue
|
||
|
|
else:
|
||
|
|
assert obj_frame['name'].startswith('O')
|
||
|
|
obj_name = obj_frame['name']
|
||
|
|
# 0: enter 1: disappear 2: update 3: unchange
|
||
|
|
frame_id = obj_frame['frame']
|
||
|
|
curr_loc = get_grid_location(obj_frame)
|
||
|
|
mind_dict = {'m1': {'fluent': 3, 'loc': None}, 'm2': {'fluent': 3, 'loc': None},
|
||
|
|
'm12': {'fluent': 3, 'loc': None},
|
||
|
|
'm21': {'fluent': 3, 'loc': None}, 'mc': {'fluent': 3, 'loc': None},
|
||
|
|
'mg': {'fluent': 3, 'loc': curr_loc}}
|
||
|
|
mind_dict['mg']['loc'] = curr_loc
|
||
|
|
if not type(obj_frame['label']) == float:
|
||
|
|
mind_labels = obj_frame['label'].split()
|
||
|
|
for mind_label in mind_labels:
|
||
|
|
if mind_label == 'in_m1' or mind_label == 'in_m2' or mind_label == 'in_m12' \
|
||
|
|
or mind_label == 'in_m21' or mind_label == 'in_mc' or mind_label == '"in_m1"' or mind_label == '"in_m2"'\
|
||
|
|
or mind_label == '"in_m12"' or mind_label == '"in_m21"' or mind_label == '"in_mc"':
|
||
|
|
mind_name = mind_label.split('_')[1].split('"')[0]
|
||
|
|
mind_dict[mind_name]['loc'] = curr_loc
|
||
|
|
else:
|
||
|
|
mind_name = mind_label.split('_')[0].split('"')
|
||
|
|
if len(mind_name) > 1:
|
||
|
|
mind_name = mind_name[1]
|
||
|
|
else:
|
||
|
|
mind_name = mind_name[0]
|
||
|
|
last_loc = obj_records[obj_name][frame_id - 1][mind_name]['loc']
|
||
|
|
mind_dict[mind_name]['loc'] = last_loc
|
||
|
|
|
||
|
|
for mind_name in mind_dict.keys():
|
||
|
|
if frame_id > 0:
|
||
|
|
curr_loc = mind_dict[mind_name]['loc']
|
||
|
|
last_loc = obj_records[obj_name][frame_id - 1][mind_name]['loc']
|
||
|
|
if last_loc is None and curr_loc is not None:
|
||
|
|
mind_dict[mind_name]['fluent'] = 0
|
||
|
|
elif last_loc is not None and curr_loc is None:
|
||
|
|
mind_dict[mind_name]['fluent'] = 1
|
||
|
|
elif not last_loc == curr_loc:
|
||
|
|
mind_dict[mind_name]['fluent'] = 2
|
||
|
|
if obj_name not in obj_records:
|
||
|
|
obj_records[obj_name] = [mind_dict]
|
||
|
|
else:
|
||
|
|
obj_records[obj_name].append(mind_dict)
|
||
|
|
|
||
|
|
with open(save_path + task.split('/')[-1].split('.')[0] + '.p', 'wb') as f:
|
||
|
|
pickle.dump(obj_records, f)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
reframe_annotation()
|
||
|
|
regenerate_annotation()
|