def test_relabel_pair():
"""
grayscale -> grayscale label remapping.
"""
old_dataroot = f'{_TEST_DIR}/test_data'
new_dataroot = f'{_TEST_DIR}/test_data'
orig_pair = ('rgb_old.jpg', 'remapping_test_data_old/label_old.png')
remapped_pair = ('rgb_new.jpg', 'remapping_test_data_new/label_new.png')
old_label_fpath = f'{old_dataroot}/{orig_pair[1]}'
create_leading_fpath_dirs(old_label_fpath)
new_label_fpath = f'{new_dataroot}/{remapped_pair[1]}'
semantic_img = np.array([[254, 0, 1], [7, 8, 9]], dtype=np.uint8)
imageio.imwrite(old_label_fpath, semantic_img)
label_mapping = {254: 253, 0: 255, 1: 0, 7: 6, 8: 7, 9: 8}
label_mapping_arr = form_label_mapping_array(label_mapping)
relabel_pair(old_dataroot,
new_dataroot,
orig_pair,
remapped_pair,
label_mapping_arr,
dataset_colors=None)
gt_mapped_img = np.array([[253, 255, 0], [6, 7, 8]], dtype=np.uint8)
remapped_img = imageio.imread(new_label_fpath)
assert np.allclose(gt_mapped_img, remapped_img)
os.remove(old_label_fpath)
os.remove(new_label_fpath)
def save_prediction_visualization(
pred_folder: str, image_path: str, image_name: str, pred: np.ndarray,
target_img: np.ndarray, id_to_class_name_map: Mapping[int,
str]) -> None:
"""
Args:
- pred_folder
- image_path
- pred
- target_img
- id_to_class_name_map
Returns:
- None
"""
image_name = Path(image_name).stem
mask_save_dir = pred_folder.replace('gray', 'rgb_mask_predictions')
grid_save_fpath = f'{mask_save_dir}/{image_name}.jpg'
rgb_img = cv2_imread_rgb(image_path)
#save_pred_vs_label_7tuple(rgb_img, pred, target_img, self.id_to_class_name_map, grid_save_fpath)
write_six_img_grid_w_embedded_names(rgb_img, pred, target_img,
id_to_class_name_map, grid_save_fpath)
overlaid_save_fpath = f'{mask_save_dir}_overlaid/{image_name}.jpg'
create_leading_fpath_dirs(overlaid_save_fpath)
frame_visualizer = Visualizer(rgb_img, metadata=None)
overlaid_img = frame_visualizer.overlay_instances(
label_map=pred, id_to_class_name_map=id_to_class_name_map)
imageio.imwrite(overlaid_save_fpath, overlaid_img)
def cv2_write_rgb(save_fpath: str, img_rgb: np.ndarray) -> None:
"""
Args:
save_fpath: string representing absolute path where image should be saved
img_rgb: (H,W,C) array representing Numpy image in RGB order
"""
img_file_type = Path(save_fpath).suffix
assert img_file_type in ['.jpg', '.png']
create_leading_fpath_dirs(save_fpath)
cv2.imwrite(save_fpath, img_rgb[:, :, ::-1])
def cv2_write_rgb(save_fpath: str, img_rgb: np.ndarray):
"""
Args:
- save_fpath
Returns:
- None
"""
img_file_type = Path(save_fpath).suffix
assert img_file_type in ['.jpg', '.png']
create_leading_fpath_dirs(save_fpath)
cv2.imwrite(save_fpath, img_rgb[:, :, ::-1])
def relabel_pair(old_dataroot: str,
new_dataroot: str,
orig_pair: Tuple[str, str],
remapped_pair: Tuple[str, str],
dname: str,
tax_converter: TaxonomyConverter,
segm_to_class: Mapping[int, int],
dataset_colors: Optional[np.ndarray] = None):
"""
No need to copy the RGB files again. We just update the label file paths.
Args:
- old_dataroot:
- new_dataroot:
- orig_pair: Tuple containing relative path to RGB image and label image
- remapped_pair: Tuple containing relative path to RGB image and label image
- label_mapping_arr:
- dataset_colors:
Returns:
- None
"""
_, orig_rel_label_fpath = orig_pair
_, remapped_rel_label_fpath = remapped_pair
old_label_fpath = f'{old_dataroot}/{orig_rel_label_fpath}'
if not os.path.exists(old_label_fpath):
print("Warning: File " + old_label_fpath + " not found!")
return
if dataset_colors is None:
label_img = imageio.imread(old_label_fpath)
else:
# remap from RGB encoded labels to 1-channel class indices
label_img_rgb = cv2_imread_rgb(old_label_fpath)
label_img = rgb_img_to_obj_cls_img(label_img_rgb, dataset_colors)
if not segm_to_class is None:
label_img_id = label_img[:, :, 0] + (label_img[:, :, 1] * 256) + (
label_img[:, :, 2] * 256**2)
label_img = np.ones(label_img.shape[:2],
dtype=np.uint8) * 255 #initialize with unlabeled
for src, dst in segm_to_class.items():
label_img[label_img_id == src] = dst
labels = torch.tensor(label_img, dtype=torch.int64)
remapped_img = tax_converter.transform_label(labels, dname)
new_label_fpath = f'{new_dataroot}/{remapped_rel_label_fpath}'
create_leading_fpath_dirs(new_label_fpath)
remapped_img = remapped_img.numpy().astype(dtype=np.uint8)
imageio.imwrite(new_label_fpath, remapped_img)
def execute_on_video(self,
max_num_frames: int = 5000,
min_resolution: int = 1080) -> None:
"""
input_file is a path to a video file.
Read frames from an RGB video file, and write overlaid
predictions into a new video file.
Args:
- None
Returns:
- None
"""
in_fname_stem = Path(self.input_file).stem
out_fname = f'{in_fname_stem}_{self.args.model_name}_universal'
out_fname += f'_scales_{self.scales_str}_base_sz_{self.args.base_size}.mp4'
output_video_fpath = f'{_ROOT}/temp_files/{out_fname}'
create_leading_fpath_dirs(output_video_fpath)
logger.info(f'Write video to {output_video_fpath}')
writer = VideoWriter(output_video_fpath)
video_fpath = '/Users/johnlamb/Downloads/sample_ffmpeg.mp4'
reader = VideoReader(self.input_file)
for frame_idx in range(reader.num_frames):
logger.info(f'On image {frame_idx}/{reader.num_frames}')
rgb_img = reader.get_frame()
if frame_idx > max_num_frames:
break
pred_label_img = self.execute_on_img(rgb_img)
# avoid blurry images by upsampling RGB before overlaying text
if np.amin(rgb_img.shape[:2]) < min_resolution:
rgb_img = resize_img_by_short_side(rgb_img, min_resolution,
'rgb')
pred_label_img = resize_img_by_short_side(
pred_label_img, min_resolution, 'label')
metadata = None
frame_visualizer = Visualizer(rgb_img, metadata)
output_img = frame_visualizer.overlay_instances(
label_map=pred_label_img,
id_to_class_name_map=self.id_to_class_name_map)
writer.add_frame(output_img)
reader.complete()
writer.complete()
def form_hstacked_imgs(img_list: List[np.ndarray],
hstack_save_fpath: str,
save_to_disk: bool = True) -> np.ndarray:
"""
Concatenate images along a horizontal axis and save them.
Accept RGB images, and convert to BGR for OpenCV to save them.
Args:
img_list: list of Numpy arrays e.g. representing different RGB visualizations of same image,
must all be of same height
hstack_save_fpath: string, representing file path
Returns:
hstack_img: Numpy array representing RGB image, containing horizontally stacked images as tiles.
"""
img_file_type = Path(hstack_save_fpath).suffix
assert img_file_type in ['.jpg', '.png']
create_leading_fpath_dirs(hstack_save_fpath)
img_h, img_w, ch = img_list[0].shape
assert ch == 3
# height and number of channels must match
assert all(img.shape[0] == img_h for img in img_list)
assert all(img.shape[2] == ch for img in img_list)
num_imgs = len(img_list)
all_widths = [img.shape[1] for img in img_list]
hstack_img = np.zeros((img_h, sum(all_widths), 3), dtype=np.uint8)
running_w = 0
for i, img in enumerate(img_list):
h, w, _ = img.shape
start = running_w
end = start + w
hstack_img[:, start:end, :] = img
running_w += w
if save_to_disk:
cv2.imwrite(hstack_save_fpath, hstack_img[:, :, ::-1])
return hstack_img
def relabel_pair(old_dataroot: str,
new_dataroot: str,
orig_pair: Tuple[str, str],
remapped_pair: Tuple[str, str],
label_mapping_arr: np.ndarray,
dataset_colors: Optional[np.ndarray] = None):
"""
No need to copy the RGB files again. We just update the label file paths.
Args:
- old_dataroot:
- new_dataroot:
- orig_pair: Tuple containing relative path to RGB image and label image
- remapped_pair: Tuple containing relative path to RGB image and label image
- label_mapping_arr:
- dataset_colors:
Returns:
- None
"""
_, orig_rel_label_fpath = orig_pair
_, remapped_rel_label_fpath = remapped_pair
old_label_fpath = f'{old_dataroot}/{orig_rel_label_fpath}'
if dataset_colors is None:
label_img = imageio.imread(old_label_fpath)
else:
# remap from RGB encoded labels to 1-channel class indices
label_img_rgb = cv2_imread_rgb(old_label_fpath)
label_img = rgb_img_to_obj_cls_img(label_img_rgb, dataset_colors)
remapped_img = map_semantic_img_fast(label_img, label_mapping_arr)
new_label_fpath = f'{new_dataroot}/{remapped_rel_label_fpath}'
create_leading_fpath_dirs(new_label_fpath)
imageio.imwrite(new_label_fpath, remapped_img)