def visualise_tracks(data_file, target_folder, id_idx=0):
if not os.path.exists(target_folder):
os.makedirs(target_folder)
f = h5py.File(data_file, "r")
#Get a colormap for the ID
tracks_n = f["tracks_n"].value[0]
colormap = make_colormap(int(tracks_n + 1))
#Save visualisations
#TODO: Get rid of the rest of the padding on the saved visualisations
#TODO: Fix the caption problem
start_count = find_start_count(list(f.keys()))
for i in range(start_count, f['frame_number'].value[0]):
frame = "frame{}".format(i)
r = f[frame]
image = r['image'].value
class_names = f["class_names"]
IDs = r['IDs'].value[:, id_idx]
colors = [colormap[int(j)] for j in IDs]
save_path = os.path.join(target_folder, "{}.jpg".format(frame))
save_instances(image,
save_path,
r['rois'],
r['masks'],
r['class_ids'],
class_names,
IDs,
colors=colors)
f.close()
def iou_track(data_file):
f = h5py.File(data_file, "r+")
start_count = find_start_count(list(f.keys()))
first_frame = "frame{}".format(start_count)
track_n = f[first_frame]['masks'].shape[2]
index_array = np.arange(1, track_n + 1)
index_array = np.stack((index_array, index_array), axis=1)
f.create_dataset("{}/IDs".format(first_frame), data=index_array)
for i in range(start_count + 1, f['frame_number'].value[0]):
frame = "frame{}".format(i)
previous_frame = "frame{}".format(i - 1)
previous_mask_n = f[previous_frame]['masks'].shape[2]
current_mask_n = f[frame]['masks'].shape[2]
index_array = np.zeros((current_mask_n, 2))
ious = np.zeros((current_mask_n, previous_mask_n))
for mask_id in range(current_mask_n):
current_box = f[frame]['rois'].value[mask_id, :]
ious[mask_id, :] = np.array([
iou(f[previous_frame]['rois'].value[previous_id, :],
current_box)
for previous_id in range(f[previous_frame]['rois'].shape[0])
])
assignments = linear_assignment(-ious)
assigned_ids = []
for assignment in assignments:
assigned_ids.append(assignment[0])
if (ious[assignment[0], assignment[1]] > 0):
index_array[assignment[0], :] = f[previous_frame]['IDs'].value[
assignment[1], 0]
else:
track_n += 1
index_array[assignment[0], :] = track_n
if (len(assignments) < ious.shape[0]):
missing_ids = [
i for i in range(current_mask_n) if i not in assigned_ids
]
for missing_id in missing_ids:
track_n += 1
index_array[missing_id, :] = track_n
f.create_dataset("{}/IDs".format(frame), data=index_array)
f["tracks_n"][0] = track_n
f.close()
def resize_data(data_file, target_file, height, width, maintain_ratio=True):
f = h5py.File(data_file, "r")
f2 = h5py.File(target_file, "w")
#TODO: change this hardcoding to automatic search
f2.create_dataset("class_names", data=f["class_names"])
f2.create_dataset("frame_number", data=f["frame_number"])
if ('transition_frame' in list(f.keys())):
f2.create_dataset("transition_frame", data=f["transition_frame"])
if ('tracks_n' in f):
f2.create_dataset("tracks_n", data=f['tracks_n'])
initial_dims = f['frame1']['image'].value.shape
if (maintain_ratio):
height, width, ratio = incorporate_ratio(initial_dims, height, width)
ratio_h = ratio
ratio_w = ratio
else:
ratio_h = height / initial_dims[0]
ratio_w = width / initial_dims[1]
start_count = find_start_count(list(f.keys()))
frame_indices = range(start_count, f['frame_number'].value[0])
for i in frame_indices:
frame = "frame{}".format(i)
for key in f[frame].keys():
if (len(f[frame][key].shape) > 2
and f[frame][key].shape[0] == initial_dims[0]
and f[frame][key].shape[1] == initial_dims[1]
and f[frame][key].shape[2] > 0):
resized_image = resize_image(f[frame][key].value, height,
width)
if (len(resized_image.shape) < 3):
resized_image = resized_image[:, :, np.newaxis]
f2.create_dataset("{}/{}".format(frame, key),
data=resized_image)
elif (key == 'rois'):
old_rois = f[frame][key].value
new_rois = np.zeros(old_rois.shape)
new_rois[:, (0, 2)] = old_rois[:, (0, 2)] * ratio_w
new_rois[:, (1, 3)] = old_rois[:, (1, 3)] * ratio_h
f2.create_dataset("{}/{}".format(frame, key),
data=new_rois.astype(int))
else:
f2.create_dataset("{}/{}".format(frame, key),
data=f[frame][key].value)
f.close()
f2.close()
def consolidate_indices(data_file, target_file=None, verbose=0):
if (target_file == None):
target_file = data_file
else:
if verbose == 1: print("Creating target file...")
#assert not os.path.exists(target_file)
copyfile(data_file, target_file)
f = h5py.File(target_file, "r+")
tracks_n = f["tracks_n"].value
start_count = find_start_count(list(f.keys()))
for i in range(start_count, f['frame_number'].value[0] - 1):
frame1 = "frame{}".format(i)
frame2 = "frame{}".format(i + 1)
IDs_1 = f[frame1]['IDs'].value
IDs_2 = f[frame2]['IDs'].value
for pair_index, id_pair_1 in enumerate(IDs_1):
if (np.any(id_pair_1 == 0)):
id_pair_1[id_pair_1 == 0] = tracks_n + 1
f[frame1]['IDs'][pair_index, :] = id_pair_1
tracks_n += 1
if (np.any(np.all(id_pair_1 == IDs_2, axis=1))):
continue
elif (np.any(id_pair_1[0] == IDs_2[:,
0])): #and not id_pair_1[0]==0
idx = np.where(IDs_2[:, 0] == id_pair_1[0])
change_to = id_pair_1[1]
change_from = IDs_2[idx[0][0], 1]
if (change_from == 0):
f[frame2]['IDs'][idx[0][0], 1] = change_to
else:
recursive_update(f, change_from, change_to, 1, i + 1,
start_count)
elif (np.any(
id_pair_1[1] == IDs_2[:, 1])): #and not id_pair_1[1] == 0
idx = np.where(IDs_2[:, 1] == id_pair_1[1])
change_to = id_pair_1[0]
change_from = IDs_2[idx[0][0], 0]
if (change_from == 0):
f[frame2]['IDs'][idx[0][0], 0] = change_to
else:
recursive_update(f, change_from, change_to, 0, i + 1,
start_count)
f["tracks_n"][0] = tracks_n
f.close()
def get_mask_stats(file, class_filtered_indices, threshold, verbose=0):
mask_list = []
start_count = find_start_count(list(file.keys()))
frame_indices = range(start_count, file['frame_number'].value[0])
for i in frame_indices:
frame = "frame{}".format(i)
for mask_idx in class_filtered_indices[i]:
mask_area = np.sum(file[frame]["masks"].value[:, :, mask_idx])
mask_list.append(mask_area)
if verbose == 2: print(mask_list)
return np.mean(mask_list), np.std(mask_list), np.percentile(
mask_list, threshold)
def score_and_pos_discard(data_file,
target_file,
scores,
positions=-1,
verbose=0):
f = h5py.File(data_file, "r")
cls_names = f["class_names"]
scores_with_ids = {}
for score_tup in scores:
scores_with_ids[np.where(
cls_names.value == np.string_(score_tup[0]))[0][0]] = score_tup[1]
if verbose == 1:
print("The min scores and ids are {}".format(scores_with_ids))
start_count = find_start_count(list(f.keys()))
frame_indices = range(start_count, f['frame_number'].value[0])
f2 = h5py.File(target_file, "w")
f2.create_dataset("class_names", data=f["class_names"])
f2.create_dataset("tracks_n", data=np.array([0]))
f2.create_dataset("frame_number", data=f["frame_number"])
if ('transition_frame' in list(f.keys())):
f2.create_dataset("transition_frame", data=f["transition_frame"])
for i in frame_indices:
frame = "frame{}".format(i)
relevant_masks = []
for mask_id in range(f[frame]["rois"].shape[0]):
if (class_and_score(f[frame], scores_with_ids, mask_id)
and correct_pos(f[frame], positions, mask_id)):
relevant_masks.append(mask_id)
f2.create_dataset("{}/class_ids".format(frame),
data=f[frame]['class_ids'].value[relevant_masks])
f2.create_dataset("{}/image".format(frame),
data=f[frame]["image"].value)
f2.create_dataset("{}/rois".format(frame),
data=f[frame]["rois"].value[relevant_masks])
f2.create_dataset("{}/scores".format(frame),
data=f[frame]["scores"].value[relevant_masks])
f2.create_dataset("{}/masks".format(frame),
data=f[frame]["masks"].value[:, :, relevant_masks])
f.close()
f2.close()
def make_dataset(data_file,
target_file,
timestep=2,
number_inputs=3,
future_time=10,
sparse_sampling=-1,
merged_data=False):
f = h5py.File(data_file, "r")
f2 = h5py.File(target_file, "w")
f2.create_dataset("origin_file", data=[np.string_(data_file)])
f2.create_dataset("timestep", data=[timestep])
f2.create_dataset("number_inputs", data=[number_inputs])
f2.create_dataset("future_time", data=[future_time])
max_frames = f['frame_number'].value[0]
start_count = find_start_count(list(f.keys()))
if (sparse_sampling > 0):
f2.create_dataset("sparse_sampling", data=[sparse_sampling])
frame_indices = range(start_count, max_frames, sparse_sampling)
else:
frame_indices = range(start_count, max_frames)
if (merged_data):
transition_frame = f['transition_frame'].value[0]
f2.create_dataset("transition_frame", data=[transition_frame])
recorded_transition_datapoint = False
datapoint_index = 0
for i in frame_indices:
frame = "frame{}".format(i)
id_list = []
for id, id_2 in f[frame]['IDs']:
#Check that the frame at t - timestep*(num_inputs-1) exists
if (i - int(timestep * (number_inputs - 1)) < 0):
continue
#Check that the future frame exists
future_frame = i + future_time
if (future_frame > max_frames - 1):
break
#Add the ids of that frame that have previous and future frames with the same id
if (check_id_consistency(f, id, i, timestep, number_inputs,
future_time)):
id_list.append(id)
#Add the datapoint
if (len(id_list) > 0):
if (merged_data):
if (i >= transition_frame
and not recorded_transition_datapoint):
f2.create_dataset("transition_datapoint",
data=[datapoint_index])
recorded_transition_datapoint = True
for id in id_list:
add_datapoint(f, f2, id, i, timestep, number_inputs,
future_time, datapoint_index)
datapoint_index += 1
f2.create_dataset("datapoints", data=[datapoint_index])
def merge_data(dataset1,dataset2,new_dataset_folder, new_dataset_name, verbose = 0):
if verbose > 0 : print("Starting data merge...")
#If the new folder does not exist, create if
if not os.path.exists(new_dataset_folder):
os.makedirs(new_dataset_folder)
#get the path of the new dataset
new_dataset = os.path.join(new_dataset_folder,'{}.hdf5'.format(new_dataset_name))
#open the files
f1 = h5py.File(dataset1, "r")
f2 = h5py.File(dataset2, "r")
f_new = h5py.File(new_dataset, "w")
if verbose > 0: print("Making metadata...")
#Put the metadata into the new file
total_frames = f1['frame_number'].value[0] + f2['frame_number'].value[0]
f_new.create_dataset("class_names", data = f1['class_names'])
f_new.create_dataset("frame_number", data = [total_frames])
f_new.create_dataset("transition_frame", data = [f1['frame_number'].value[0]])
count = 0
if verbose > 0: print("Copying over frames from {}".format(dataset1))
#Iterate over the frames of the the first video
start_count = find_start_count(list(f1.keys()))
frame_indices = range(start_count, f1['frame_number'].value[0])
for i in frame_indices:
frame = 'frame{}'.format(i)
for k, v in f1[frame].items():
f_new.create_dataset("frame{}/{}".format(count,k), data=v)
count+=1
if verbose > 0: print("Asserting...")
#Check all the keys are correct
aset = set(f_new.keys())
aset.remove('transition_frame')
assert set(f1.keys()) == aset
assert set(f1['frame4'].keys()) == set(f_new['frame4'].keys())
#Check that some of the values are correct
rand_indices = np.random.randint(0,f1['frame_number'].value[0]-1,5)
for idx in rand_indices:
frame = 'frame{}'.format(idx)
for k,v in f_new[frame].items():
assert np.all(v.value == f1[frame][k].value)
if verbose > 0: print("Copying over frames from {}".format(dataset2))
#Iterate over the frames of the second video
start_count = find_start_count(list(f2.keys()))
frame_indices = range(start_count, f2['frame_number'].value[0])
for i in frame_indices:
frame = 'frame{}'.format(i)
for k, v in f2[frame].items():
f_new.create_dataset("frame{}/{}".format(count,k), data=v)
count+=1
if verbose > 0: print("Asserting...")
#Check that the keys are correct
assert total_frames == count
assert set(f2['frame{}'.format(f2['frame_number'].value[0] - 5)])==set(f_new['frame{}'.format(total_frames -5)])
#Check that some of the values are correct
rand_indices = np.random.randint(0, f2['frame_number'].value[0] - 1, 5)
for idx in rand_indices:
frame_old = 'frame{}'.format(idx)
frame_new = 'frame{}'.format(idx + f1['frame_number'].value[0])
for k, v in f_new[frame_new].items():
assert np.all(v.value == f2[frame_old][k].value)
f1.close()
f2.close()
f_new.close()
def class_and_size_discard(data_file,
target_file,
masks_to_keep=['car'],
small_threshold=-1,
global_stats=True,
verbose=0):
f = h5py.File(data_file, "r")
cls_names = f["class_names"]
mask_bstrings = [np.string_(j) for j in masks_to_keep]
mask_ids = [
np.where(cls_names.value == name)[0][0] for name in mask_bstrings
]
if verbose == 1: print("The mask_ids to look for {}".format(mask_ids))
start_count = find_start_count(list(f.keys()))
frame_indices = range(start_count, f['frame_number'].value[0])
class_filtered_indices = []
frame_percentile = []
if (start_count == 1):
class_filtered_indices.append(None)
frame_percentile.append(None)
for i in frame_indices:
frame = "frame{}".format(i)
class_indices = np.where(np.isin(f[frame]['class_ids'].value,
mask_ids))[0]
class_filtered_indices.append(class_indices)
if (not global_stats and small_threshold > 0):
mean, std, percentile = get_mask_stats_local(
f[frame], class_indices, small_threshold, verbose)
frame_percentile.append(percentile)
if (small_threshold > 0 and global_stats):
mean, std, percentile = get_mask_stats(f, class_filtered_indices,
small_threshold, verbose)
frame_percentile = [percentile] * len(class_filtered_indices)
if verbose == 1:
print("mean area : {}, std aread : {}, percentile {}".format(
mean, std, percentile))
f2 = h5py.File(target_file, "w")
f2.create_dataset("class_names", data=f["class_names"])
f2.create_dataset("tracks_n", data=np.array([0]))
f2.create_dataset("frame_number", data=f["frame_number"])
if ('transition_frame' in list(f.keys())):
f2.create_dataset("transition_frame", data=f["transition_frame"])
for i in frame_indices:
frame = "frame{}".format(i)
if (small_threshold > 0):
relevant_masks = []
for mask_id in class_filtered_indices[i]:
if (np.sum(f[frame]["masks"].value[:, :, mask_id]) >
frame_percentile[i]):
relevant_masks.append(mask_id)
elif (verbose == 1):
print("Discarding mask with area {}".format(
np.sum(f[frame]["masks"].value[:, :, mask_id])))
else:
relevant_masks = class_filtered_indices[i]
f2.create_dataset("{}/class_ids".format(frame),
data=f[frame]['class_ids'].value[relevant_masks])
f2.create_dataset("{}/image".format(frame),
data=f[frame]["image"].value)
f2.create_dataset("{}/rois".format(frame),
data=f[frame]["rois"].value[relevant_masks])
f2.create_dataset("{}/scores".format(frame),
data=f[frame]["scores"].value[relevant_masks])
f2.create_dataset("{}/masks".format(frame),
data=f[frame]["masks"].value[:, :, relevant_masks])
f.close()
f2.close()
def track(data_file, reverse=False, verbose=0):
if (verbose == 1):
print("Opening File...")
f = h5py.File(data_file, "r+")
mot_tracker = Sort()
tracks_n = f["tracks_n"].value[0]
start_count = find_start_count(list(f.keys()))
if (not reverse):
frame_indices = range(start_count, f['frame_number'].value[0])
else:
frame_indices = reversed(range(start_count,
f['frame_number'].value[0]))
if (verbose == 1):
print("Starting loop...")
for i in frame_indices:
frame = "frame{}".format(i)
bbox_handle = f[frame]['rois']
detection = bbox_handle.value
scores = f[frame]['scores'].value
number_of_masks = scores.shape[0]
detection_with_scores = np.hstack(
(detection, np.reshape(scores, (-1, 1))))
if (verbose == 1):
print("detections with scores:")
print(detection_with_scores)
track_bbs_ids = mot_tracker.update(detection_with_scores)
if (verbose == 1):
print("tracked bbs:")
print(track_bbs_ids)
# Associate the track_BBs with the original bbs
# for each of the track bbs
# find the nearest neighbour in the original detections
# associate the ID with the index of the original detection
index_array = np.zeros(number_of_masks)
if verbose == 1: print("number of masks {}".format(number_of_masks))
for track in track_bbs_ids:
nn_index = find_nn(track[:-1], detection)
index_array[nn_index] = track[-1]
if (verbose == 1):
print("The index array is")
print(index_array)
max_idx = np.amax(index_array) if number_of_masks > 0 else 0
if (max_idx > tracks_n):
tracks_n = max_idx
ID_dataset_key = "{}/IDs".format(frame)
if (ID_dataset_key in f):
f[ID_dataset_key][:, 1] = index_array
else:
f.create_dataset(ID_dataset_key, (index_array.shape[0], 2))
f[ID_dataset_key][:, 0] = index_array
f["tracks_n"][0] = tracks_n
KalmanBoxTracker.count = 0
f.close()