96 lines
4.9 KiB
Python
96 lines
4.9 KiB
Python
|
import os
|
||
|
|
import pandas as pd
|
||
|
|
import pickle
|
||
|
|
from tqdm import tqdm
|
||
|
|
|
||
|
|
|
||
|
|
def check_append(obj_name, m1, mind_name, obj_frame, flags, label):
|
||
|
|
if label:
|
||
|
|
if not obj_name in m1:
|
||
|
|
m1[obj_name] = []
|
||
|
|
m1[obj_name].append(
|
||
|
|
[obj_frame.frame, [obj_frame.x_min, obj_frame.y_min, obj_frame.x_max, obj_frame.y_max], 0])
|
||
|
|
flags[mind_name] = 1
|
||
|
|
elif not flags[mind_name]:
|
||
|
|
m1[obj_name].append(
|
||
|
|
[obj_frame.frame, [obj_frame.x_min, obj_frame.y_min, obj_frame.x_max, obj_frame.y_max], 0])
|
||
|
|
flags[mind_name] = 1
|
||
|
|
else: # false belief
|
||
|
|
if obj_name in m1:
|
||
|
|
if flags[mind_name]:
|
||
|
|
m1[obj_name].append(
|
||
|
|
[obj_frame.frame, [obj_frame.x_min, obj_frame.y_min, obj_frame.x_max, obj_frame.y_max], 1])
|
||
|
|
flags[mind_name] = 0
|
||
|
|
return flags, m1
|
||
|
|
|
||
|
|
|
||
|
|
def store_mind_set(clip, annotation_path, save_path):
|
||
|
|
if not os.path.exists(save_path):
|
||
|
|
os.makedirs(save_path)
|
||
|
|
annt = pd.read_csv(annotation_path + clip, sep=",", header=None)
|
||
|
|
annt.columns = ["obj_id", "x_min", "y_min", "x_max", "y_max", "frame", "lost", "occluded", "generated", "name",
|
||
|
|
"label"]
|
||
|
|
obj_names = annt.name.unique()
|
||
|
|
m1, m2, m12, m21, mc = {}, {}, {}, {}, {}
|
||
|
|
flags = {'m1':0, 'm2':0, 'm12':0, 'm21':0, 'mc':0}
|
||
|
|
for obj_name in obj_names:
|
||
|
|
if obj_name == 'P1' or obj_name == 'P2':
|
||
|
|
continue
|
||
|
|
obj_frames = annt.loc[annt.name == obj_name]
|
||
|
|
for index, obj_frame in obj_frames.iterrows():
|
||
|
|
if type(obj_frame.label) == float:
|
||
|
|
continue
|
||
|
|
labels = obj_frame.label.split()
|
||
|
|
for label in labels:
|
||
|
|
if label == 'in_m1' or label == '"in_m1"':
|
||
|
|
flags, m1 = check_append(obj_name, m1, 'm1', obj_frame, flags, 1)
|
||
|
|
elif label == 'in_m2' or label == '"in_m2"':
|
||
|
|
flags, m2 = check_append(obj_name, m2, 'm2', obj_frame, flags, 1)
|
||
|
|
elif label == 'in_m12'or label == '"in_m12"':
|
||
|
|
flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 1)
|
||
|
|
elif label == 'in_m21' or label == '"in_m21"':
|
||
|
|
flags, m21 = check_append(obj_name, m21, 'm21', obj_frame, flags, 1)
|
||
|
|
elif label == 'in_mc'or label == '"in_mc"':
|
||
|
|
flags, mc = check_append(obj_name, mc, 'mc', obj_frame, flags, 1)
|
||
|
|
elif label == 'm1_false' or label == '"m1_false"':
|
||
|
|
flags, m1 = check_append(obj_name, m1, 'm1', obj_frame, flags, 0)
|
||
|
|
flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 0)
|
||
|
|
flags, m21 = check_append(obj_name, m21, 'm21', obj_frame, flags, 0)
|
||
|
|
false_belief = 'm1_false'
|
||
|
|
with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
|
||
|
|
file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
|
||
|
|
elif label == 'm2_false' or label == '"m2_false"':
|
||
|
|
flags, m2 = check_append(obj_name, m2, 'm2', obj_frame, flags, 0)
|
||
|
|
flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 0)
|
||
|
|
flags, m21 = check_append(obj_name, m21, 'm21', obj_frame, flags, 0)
|
||
|
|
false_belief = 'm2_false'
|
||
|
|
with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
|
||
|
|
file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
|
||
|
|
elif label == 'm12_false' or label == '"m12_false"':
|
||
|
|
flags, m12 = check_append(obj_name, m12, 'm12', obj_frame, flags, 0)
|
||
|
|
flags, mc = check_append(obj_name, mc, 'mc', obj_frame, flags, 0)
|
||
|
|
false_belief = 'm12_false'
|
||
|
|
with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
|
||
|
|
file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
|
||
|
|
elif label == 'm21_false' or label == '"m21_false"':
|
||
|
|
flags, m21 = check_append(obj_name, m2, 'm21', obj_frame, flags, 0)
|
||
|
|
flags, mc = check_append(obj_name, mc, 'mc', obj_frame, flags, 0)
|
||
|
|
false_belief = 'm21_false'
|
||
|
|
with open(save_path + clip.split('.')[0] + '.txt', "a") as file:
|
||
|
|
file.write(f"{obj_frame['frame']},{obj_name},{false_belief}\n")
|
||
|
|
# print('m1', m1)
|
||
|
|
# print('m2', m2)
|
||
|
|
# print('m12', m12)
|
||
|
|
# print('m21', m21)
|
||
|
|
# print('mc', mc)
|
||
|
|
#with open(save_path + clip.split('.')[0] + '.p', 'wb') as f:
|
||
|
|
# pickle.dump([m1, m2, m12, m21, mc], f)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
|
||
|
|
annotation_path = '/scratch/bortoletto/data/tbd/reformat_annotation/'
|
||
|
|
save_path = '/scratch/bortoletto/data/tbd/store_mind_set/'
|
||
|
|
|
||
|
|
for clip in tqdm(os.listdir(annotation_path), desc="Processing videos", unit="item"):
|
||
|
|
store_mind_set(clip, annotation_path, save_path)
|