def __add_missing_boxes_to_mask2(mask1: np.ndarray, mask2: np.ndarray, distance_threshold=10):
mask1 = mask1.copy()
mask2 = mask2.copy()
mask1_rps = get_region_props(mask1)
mask2_rps = get_region_props(mask2)
missing_boxes = []
# if there are boxes in mask 1 but there is no box in mask 2
# add all the boxes of mask 1 to mask 2
if len(mask1_rps) > 0 and len(mask2_rps) == 0:
missing_boxes.extend(mask1_rps)
# if there are some boxes in mask 1 and mask 2
# If these are not present in mask2 these will be copied to that
elif len(mask1_rps) > 0 and len(mask2_rps) > 0:
distance_matrix = get_distance_matrix(mask1_rps, mask2_rps)
for i, dm in enumerate(distance_matrix):
min_distance = np.min(dm)
# if current box of mask 1 doesn't exist in mask 2
# add this mask in mask 2
if min_distance > distance_threshold:
missing_boxes.append(mask1_rps[i])
for missing_box in missing_boxes:
foreground_number = int(np.max(mask1))
y1, x1, y2, x2 = missing_box.bbox
mask2[y1:y2, x1:x2] = foreground_number
return mask2
def main_process(params):
frame_number = params['frame_number']
video_masks = params['video_masks']
frame_paths = params['frame_paths']
left_frame_number = frame_number - window_of_search
right_frame_number = frame_number + window_of_search
left_mask_path = video_masks[left_frame_number]
right_mask_path = video_masks[right_frame_number]
middle_mask_path = video_masks[frame_number]
left_frame_index = get_index_of_frame(left_mask_path, frame_paths)
right_frame_index = get_index_of_frame(right_mask_path, frame_paths)
middle_frame_index = get_index_of_frame(middle_mask_path, frame_paths)
if left_frame_index is None or right_frame_index is None or middle_frame_index is None:
print(
f'Left, middle or right frame was not found {os.path.basename(middle_mask_path)}\nContinuing by neglecting it.'
)
return ''
# from left index to before middle index
previous_frames = frame_paths[left_frame_index:middle_frame_index]
previous_frames = previous_frames[-1::-1]
# from next of middle index to right index
next_frames = frame_paths[middle_frame_index + 1:right_frame_index + 1]
middle_frame = cv2.imread(frame_paths[middle_frame_index])
middle_mask = cv2.imread(middle_mask_path)
middle_mask_rps = get_region_props(middle_mask)
for middle_mask_rp in middle_mask_rps:
y1, x1, y2, x2 = middle_mask_rp.bbox
previous_frame_presence = track_in_frames(middle_frame,
(y1, x1, y2, x2),
previous_frames,
masks_folder)
next_frame_presence = track_in_frames(middle_frame, (y1, x1, y2, x2),
next_frames, masks_folder)
# if box is present in either previous or next frame for number of times
if previous_frame_presence >= should_present_in_frames or next_frame_presence >= should_present_in_frames:
is_tp = True
else:
is_tp = False
if not is_tp:
y1, x1, y2, x2 = get_safe_extended_pixels(middle_mask.shape[0],
middle_mask.shape[1], 1,
y1, x1, y2, x2)
middle_mask[y1:y2, x1:x2] = 0
middle_mask_filename = os.path.basename(middle_mask_path)
middle_mask_filename_wo_ext, _ = os.path.splitext(middle_mask_filename)
mask_output_path = os.path.join(masks_output_folder,
f'{middle_mask_filename_wo_ext}.png')
cv2.imwrite(mask_output_path, middle_mask)
def merge_boxes_using_distance(mask1: np.ndarray, mask2: np.ndarray, distance_threshold=30,
log_file_path=None, filename=None) -> np.ndarray:
mask1 = mask1.copy()
mask2 = mask2.copy()
if len(mask1.shape) == 3 and mask1.shape[2] == 3:
mask1 = cv2.cvtColor(mask1, cv2.COLOR_BGR2GRAY)
if len(mask2.shape) == 3 and mask2.shape[2] == 3:
mask2 = cv2.cvtColor(mask2, cv2.COLOR_BGR2GRAY)
mask1_region_props = get_region_props(mask1)
mask2_region_props = get_region_props(mask2)
new_mask = mask1.copy()
if len(mask1_region_props) == 0:
for mask2_region_prop in mask2_region_props:
y1, x1, y2, x2 = mask2_region_prop.bbox
new_mask[y1:y2, x1:x2] = 255
if log_file_path is not None and filename is not None:
write_log_file(log_file_path, filename, y1, x1, y2, x2)
else:
distance_matrix = np.zeros((len(mask2_region_props), len(mask1_region_props)), dtype=np.float)
for mask2_region_prop_index, mask2_region_prop in enumerate(mask2_region_props):
y2, x2 = mask2_region_prop.centroid
for mask1_region_prop_index, mask1_region_prop in enumerate(mask1_region_props):
y1, x1 = mask1_region_prop.centroid
distance = calculate_distance(y1, x1, y2, x2)
distance_matrix[mask2_region_prop_index, mask1_region_prop_index] = distance
for mask2_region_prop_index, box_distances in enumerate(distance_matrix):
# minimum distance between current box with the boxes mask 1
min_distance = float(np.min(box_distances))
if min_distance > distance_threshold:
region_prop = mask2_region_props[mask2_region_prop_index]
y1, x1, y2, x2 = region_prop.bbox
new_mask[y1:y2, x1:x2] = 255
if log_file_path is not None and filename is not None:
write_log_file(log_file_path, filename, y1, x1, y2, x2)
return new_mask
def remove_boxes_less_than_threshold(masks_folder: str, threshold=50):
mask_paths = glob(os.path.join(masks_folder, '*'))
for mask_path in tqdm(mask_paths, desc="Removing small boxes"):
mask = cv.imread(mask_path)
rps = get_region_props(mask)
new_mask = np.zeros(mask.shape[0:2], dtype=mask.dtype)
for rp in rps:
y1, x1, y2, x2 = rp.bbox
if rp.bbox_area >= threshold:
new_mask[y1:y2, x1:x2] = 255
cv.imwrite(mask_path, new_mask)
def merge_boxes_using_iou(mask1: np.ndarray, mask2: np.ndarray, iou_threshold=0.1) -> np.ndarray:
mask1 = mask1.copy()
mask2 = mask2.copy()
if len(mask1.shape) == 3 and mask1.shape[2] == 3:
mask1 = cv2.cvtColor(mask1, cv2.COLOR_BGR2GRAY)
if len(mask2.shape) == 3 and mask2.shape[2] == 3:
mask2 = cv2.cvtColor(mask2, cv2.COLOR_BGR2GRAY)
mask1_region_props = get_region_props(mask1)
mask2_region_props = get_region_props(mask2)
new_mask = mask1.copy()
if len(mask1_region_props) == 0:
for mask2_region_prop in mask2_region_props:
y1, x1, y2, x2 = mask2_region_prop.bbox
new_mask[y1:y2, x1:x2] = 255
else:
iou_matrix = np.zeros((len(mask2_region_props), len(mask1_region_props)), dtype=np.float)
for mask2_region_prop_index, mask2_region_prop in enumerate(mask2_region_props):
mask2_y1, mask2_x1, mask2_y2, mask2_x2 = mask2_region_prop.bbox
for mask1_region_prop_index, mask1_region_prop in enumerate(mask1_region_props):
mask1_y1, mask1_x1, mask1_y2, mask1_x2 = mask1_region_prop.bbox
box1 = (mask1_x1, mask1_y1, mask1_x2, mask1_y2)
box2 = (mask2_x1, mask2_y1, mask2_x2, mask2_y2)
iou = bb_intersection_over_union(box1, box2)
iou_matrix[mask2_region_prop_index, mask1_region_prop_index] = iou
for mask2_region_prop_index, iou_values in enumerate(iou_matrix):
# maximum iou between current box with the boxes of mask 1
max_iou = float(np.max(iou_values))
if max_iou < iou_threshold:
region_prop = mask2_region_props[mask2_region_prop_index]
y1, x1, y2, x2 = region_prop.bbox
new_mask[y1:y2, x1:x2] = 255
return new_mask
def region_prop_masks(masks_folder: str):
mask_paths = glob(os.path.join(masks_folder, '*'))
for mask_path in tqdm(mask_paths, desc="Region proping masks"):
mask = cv.imread(mask_path)
rps = get_region_props(mask)
new_mask = np.zeros(mask.shape, dtype=mask.dtype)
foreground_pixel = np.max(mask)
for rp in rps:
y1, x1, y2, x2 = rp.bbox
new_mask[y1:y2, x1:x2] = foreground_pixel
cv.imwrite(mask_path, new_mask)
def main():
os.makedirs(output_folder, exist_ok=True)
mask_paths = glob(os.path.join(masks_folder, '*'))
for mask_path in tqdm(mask_paths):
filename = os.path.basename(mask_path)
mask = cv2.imread(mask_path)
rps = get_region_props(mask)
new_mask = np.zeros(mask.shape, mask.dtype)
for rp in rps:
if min_box_threshold <= rp.bbox_area < max_box_threshold:
y1, x1, y2, x2 = rp.bbox
new_mask[y1:y2, x1:x2] = 255
output_path = os.path.join(output_folder, filename)
cv2.imwrite(output_path, new_mask)
def is_box_present_in_mask(box, mask, distance_threshold):
y1, x1, y2, x2 = box
cx = x1 + (x2 - x1) / 2
cy = y1 + (y2 - y1) / 2
distances = []
mask_rps = get_region_props(mask)
for mask_rp in mask_rps:
cy2, cx2 = mask_rp.centroid
distance = calculate_distance(cy, cx, cy2, cx2)
distances.append(distance)
if len(distances) == 0:
return False
else:
min_distance = min(distances)
if min_distance <= distance_threshold:
return True
else:
return False
def main():
args = get_args()
frames_folder = args.frames_folder
labels_mask_folder = args.labels_mask_folder
output_folder = args.output_folder
save_boxes_as_json = args.save_boxes_as_json
output_json_folder = output_folder + '_json'
os.makedirs(output_folder, exist_ok=True)
if save_boxes_as_json == 1:
os.makedirs(output_json_folder, exist_ok=True)
frame_paths = glob(path.join(frames_folder, '*.png'))
frame_paths.sort()
for frame_path in tqdm(frame_paths):
filename = path.basename(frame_path)
filename_wo_ext, _ = path.splitext(filename)
labels_mask_path = path.join(labels_mask_folder, filename)
frame = cv2.imread(frame_path)
labels_mask = cv2.imread(labels_mask_path)
labels_mask = cv2.cvtColor(labels_mask, cv2.COLOR_BGR2GRAY)
crf_mask = np.zeros(labels_mask.shape, dtype=np.uint8)
crf_boxes = []
region_props = get_region_props(labels_mask)
for region_prop in region_props:
__labels_mask = np.zeros(labels_mask.shape, labels_mask.dtype)
y1, x1, y2, x2 = region_prop.bbox
__labels_mask[y1:y2, x1:x2] = 255
y1, x1, y2, x2 = get_safe_extended_pixels(labels_mask.shape[0], labels_mask.shape[1], patch_extra_pixels,
y1, x1, y2, x2)
crf_mask_patch = calculate_crf_on_labels(frame[y1:y2, x1:x2], __labels_mask[y1:y2, x1:x2])
temp_mask = np.zeros(labels_mask.shape, dtype=np.uint8)
temp_mask[y1:y2, x1:x2] = crf_mask_patch
# extending foreground mask before saving
temp_mask = expand_labels(temp_mask, 3)
# to make sure only foreground pixels are copied and not the background one
crf_mask[temp_mask > 0] = temp_mask[temp_mask > 0]
# for saving as json file
temp_rps = get_region_props(temp_mask)
for temp_rp in temp_rps:
temp_y1, temp_x1, temp_y2, temp_x2 = temp_rp.bbox
crf_boxes.append([temp_y1, temp_x1, temp_y2, temp_x2])
crf_mask = convert_segmented_area_to_bounding_box(crf_mask)
output_file_path = path.join(output_folder, filename)
cv2.imwrite(output_file_path, crf_mask)
if save_boxes_as_json == 1:
output_json_file = path.join(output_json_folder, f'{filename_wo_ext}.json')
with open(output_json_file, 'w') as f:
json.dump(crf_boxes, f)
def get_tracked_volumes(frames_list: List[np.ndarray],
masks_list: List[np.ndarray],
key_frame_number: int,
detection_masks_list: List[np.ndarray] = None):
key_frame = frames_list[key_frame_number]
# from frame 0 to less than key frame
previous_frames = frames_list[0:key_frame_number]
# reverse frames
previous_frames = previous_frames[-1::-1]
# from keyframe+1 to end
next_frames = frames_list[key_frame_number + 1:]
final_volumes = []
# get boxes for key frames
boxes = get_region_props(masks_list[key_frame_number])
if detection_masks_list is not None:
boxes.extend(get_region_props(detection_masks_list[key_frame_number]))
for box in boxes:
y1, x1, y2, x2 = box.bbox
w = x2 - x1
h = y2 - y1
if w < 10:
w = 10
if h < 10:
h = 10
tracker = cv.TrackerCSRT_create()
try:
tracker.init(key_frame, (x1, y1, w, h))
except:
print('Box ignored because tracker could not initialize')
continue
key_frame_patch_location = get_safe_extended_pixels(
key_frame.shape[0], key_frame.shape[1], extra_pixels, y1, x1, y2,
x2)
extended_height = key_frame_patch_location[
2] - key_frame_patch_location[0]
extended_width = key_frame_patch_location[
3] - key_frame_patch_location[1]
key_frame_patch = key_frame[
key_frame_patch_location[0]:key_frame_patch_location[2],
key_frame_patch_location[1]:key_frame_patch_location[3]]
previous_patches = get_patches_by_tracker_extend_version(
tracker, previous_frames, extended_height, extended_width)
previous_patches = previous_patches[-1::-1]
next_patches = get_patches_by_tracker_extend_version(
tracker, next_frames, extended_height, extended_width)
final_volume = []
final_volume.extend(previous_patches)
final_volume.append(key_frame_patch)
final_volume.extend(next_patches)
final_volumes.append({
'location': key_frame_patch_location,
'volume': final_volume
})
return final_volumes