def save_classnames_in_image_sufficientpx(
rgb_img: np.ndarray,
label_img: np.ndarray,
id_to_class_name_map: Mapping[int,str],
font_color = (0,0,0),
save_to_disk: bool = False,
save_fpath: str = '',
min_conncomp_px: int = 4000,
font_scale: int = 1
):
"""
Write a classname over each connected component of a label
map as long as the connected component has a sufficiently
large number of pixels (specified as argument).
Args:
- rgb_img: Numpy array (H,W,3) representing RGB image
- label_img: Numpy array (H,W) representing label map
- id_to_class_name_map: mapping from class ID to classname
- font_color: 3-tuple representing RGB font color
- save_to_disk: whether to save image to disk
- save_fpath: absolute file path
- min_conncomp_px: minimum number of pixels to justify
placing a text label over connected component
- font_scale: scale of font text
Returns:
- rgb_img: Numpy array (H,W,3) with embedded classanmes
"""
H, W, C = rgb_img.shape
class_to_conncomps_dict = scipy_conn_comp(label_img)
for class_idx, conncomps_list in class_to_conncomps_dict.items():
for conncomp in conncomps_list:
if conncomp.sum() < min_conncomp_px:
continue
text = id_to_class_name_map[class_idx]
y,x = get_mean_mask_location(conncomp)
x -= 55 # move the text so approx. centered over mask.
x = max(0,x)
x = min(W-1,x)
# jitter location if nonconvex object mean not within its mask
if conncomp[y,x] != 1:
x,y = search_jittered_location_in_mask(x,y,conncomp)
# print(f'Class idx: {class_idx}: (x,y)=({x},{y})')
rgb_img = add_text_cv2(
rgb_img,
text,
coords_to_plot_at=(x,y),
font_color=font_color,
font_scale=font_scale,
thickness=2)
if save_to_disk:
cv2_write_rgb(save_fpath, rgb_img)
return rgb_img
def test_scipy_conn_comp():
""" Make sure we can recover a dictionary of binary masks for each conn. component"""
# toy semantic label map / label image
img = np.array(
[
[1,1,2,3],
[1,4,5,3],
[0,0,1,1]
])
class_to_conncomps_dict = scipy_conn_comp(img)
print(class_to_conncomps_dict)
gt_class_to_conncomps_dict = {
0: [
np.array([ [0, 0, 0, 0],
[0, 0, 0, 0],
[1, 1, 0, 0]], dtype=np.uint8)
],
1: [
np.array([[1, 1, 0, 0],
[1, 0, 0, 0],
[0, 0, 0, 0]], dtype=np.uint8),
np.array([[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 1, 1]], dtype=np.uint8)
],
2: [
np.array([[0, 0, 1, 0],
[0, 0, 0, 0],
[0, 0, 0, 0]], dtype=np.uint8)
],
3: [
np.array([[0, 0, 0, 1],
[0, 0, 0, 1],
[0, 0, 0, 0]], dtype=np.uint8)
],
4: [
np.array([[0, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 0, 0]], dtype=np.uint8)
],
5: [
np.array([ [0, 0, 0, 0],
[0, 0, 1, 0],
[0, 0, 0, 0]], dtype=np.uint8)
]
}
for class_idx, conncomps_list in class_to_conncomps_dict.items():
gt_conncomps_list = gt_class_to_conncomps_dict[class_idx]
for conncomp, gtconncomp in zip(conncomps_list, gt_conncomps_list):
assert np.allclose(conncomp, gtconncomp)
def save_classnames_in_image_sufficientpx(rgb_img,
label_img,
id_to_class_name_map,
font_color=(0, 0, 0),
save_to_disk: bool = False,
save_fpath: str = '',
min_conncomp_px: int = 4000):
"""
Args:
- rgb_img
- label_img
- id_to_class_name_map: Mapping[int,str]
Returns:
-
"""
H, W, C = rgb_img.shape
class_to_conncomps_dict = scipy_conn_comp(label_img)
for class_idx, conncomps_list in class_to_conncomps_dict.items():
for conncomp in conncomps_list:
if conncomp.sum() < min_conncomp_px:
continue
text = id_to_class_name_map[class_idx]
y, x = get_mean_mask_location(conncomp)
x -= 55 # move the text so approx. centered over mask.
x = max(0, x)
x = min(W - 1, x)
# jitter location if nonconvex object mean not within its mask
if conncomp[y, x] != 1:
x, y = search_jittered_location_in_mask(x, y, conncomp)
# print(f'Class idx: {class_idx}: (x,y)=({x},{y})')
rgb_img = add_text_cv2(rgb_img,
text,
coords_to_plot_at=(x, y),
font_color=font_color,
font_scale=1,
thickness=2)
if save_to_disk:
cv2_write_rgb(save_fpath, rgb_img)
return rgb_img
def save_classnames_in_image_maxcardinality(
rgb_img,
label_img,
id_to_class_name_map,
font_color = (0,0,0),
save_to_disk: bool = False,
save_fpath: str = ''
) -> np.ndarray:
"""
Args:
rgb_img
label_img
id_to_class_name_map: Mapping[int,str]
Returns:
rgb_img
"""
H, W, C = rgb_img.shape
class_to_conncomps_dict = scipy_conn_comp(label_img)
for class_idx, conncomps_list in class_to_conncomps_dict.items():
mask_idx = find_max_cardinality_mask(conncomps_list)
maxsz_conncomp = conncomps_list[mask_idx]
text = id_to_class_name_map[class_idx]
y,x = get_mean_mask_location(maxsz_conncomp)
x -= 55
x = max(0,x)
x = min(W-1,x)
# print(f'Class idx: {class_idx}: (x,y)=({x},{y})')
rgb_img = add_text_cv2(
rgb_img,
text,
coords_to_plot_at=(x,y),
font_color=font_color,
font_scale = 1,
thickness=2)
if save_to_disk:
cv2_write_rgb(save_fpath, rgb_img)
return rgb_img
