def test_calculate_iou_df(self):
from supermariopy.denseposelib import calculate_iou_df
A = np.ones((10, 10), dtype=np.int)
B = np.ones((10, 10), dtype=np.int)
B[:5, :5] = 0
B[5:, 5:] = 1
B[5:, :5] = 2
predicted = np.stack([A] * 10, axis=0)
target = np.stack([B] * 10, axis=0)
label_names = ["zeros", "ones", "twos", "threes"]
df = calculate_iou_df(predicted, target, label_names)
assert np.allclose(df.zeros, np.zeros((10, )))
assert np.allclose(df.ones, np.ones((10, )) * 0.5)
assert np.allclose(df.twos, np.zeros((10, )))
assert np.allclose(df.threes, np.ones((10, )) * -1.0)
def test_calculate_overall_iou_from_df(self):
from supermariopy.denseposelib import (
calculate_overall_iou_from_df,
calculate_iou_df,
)
A = np.ones((10, 10), dtype=np.int)
B = np.ones((10, 10), dtype=np.int)
B[:5, :5] = 0
B[5:, 5:] = 1
B[5:, :5] = 2
predicted = np.stack([A] * 10, axis=0)
target = np.stack([B] * 10, axis=0)
label_names = ["zeros", "ones", "twos", "threes"]
df = calculate_iou_df(predicted, target, label_names)
df_mean = calculate_overall_iou_from_df(df)
print(df_mean)
np.testing.assert_almost_equal(df_mean["overall"], np.array([0.5 / 3]))
def main(infer_dir, output_folder, run_crf_config, n_processes):
os.makedirs(output_folder, exist_ok=True)
with open(run_crf_config, "r") as f:
config = yaml.load(f)
segmentation_algorithm_args = config["segmentation_algorithm_args"]
npz_files = glob.glob(os.path.join(infer_dir, "*.npz"))
npz_files = sorted(npz_files)
print("Using files :")
print(npz_files)
segmentation_algorithm = crf.SegmentationFromKeypoints(
**segmentation_algorithm_args)
data = []
with closing(Pool(n_processes)) as p:
for outputs in tqdm.tqdm(p.imap(load_npz, npz_files)):
data.append(outputs)
data = list_of_dicts2dict_of_lists(data)
data = {k: np.concatenate(data[k]) for k in ["image", "gauss_yx"]}
data["gauss_yx"] = data["gauss_yx"][..., ::-1]
process_func = functools.partial(
process_batches, **{
"segmentation_algorithm": segmentation_algorithm,
})
tuples = list(
zip(np.array_split(data["image"], n_processes, 0),
np.array_split(data["gauss_yx"], n_processes, 0)))
processed_data = []
with closing(Pool(n_processes)) as p:
for outputs in tqdm.tqdm(p.imap(process_func, tuples)):
processed_data.append(outputs)
labels = np.concatenate([p["labels"] for p in processed_data], 0)
labels_rgb = np.concatenate([p["labels_rgb"] for p in processed_data], 0)
heatmaps = np.concatenate([p["heatmaps"] for p in processed_data], 0)
ims_with_keypoints = np.concatenate(
[p["ims_with_keypoints"] for p in processed_data], 0)
target_dir = os.path.join(output_folder, "01_keypoints")
os.makedirs(target_dir, exist_ok=True)
write_rgb(ims_with_keypoints, target_dir, n_processes)
target_dir = os.path.join(output_folder, "02_heatmaps")
os.makedirs(target_dir, exist_ok=True)
write_rgb(heatmaps, target_dir, n_processes)
target_dir = os.path.join(output_folder, "03_labels_rgb")
os.makedirs(target_dir, exist_ok=True)
write_rgb(labels_rgb, target_dir, n_processes)
densepose_csv_path = config["densepose_csv_path"]
data_root = config["data_root"]
fname_col = config["data_fname_col"]
iuv_files = get_iuv_files(densepose_csv_path, data_root, len(labels),
fname_col)
iuvs = [cv2.imread(x, -1) for x in iuv_files]
iuvs = [
denseposelib.resize_labels(i[..., 0], labels.shape[1:]) for i in iuvs
]
iuvs = np.stack(iuvs, axis=0)
dp_semantic_remap_dict = config["dp_semantic_remap_dict"]
dp_new_part_list = sorted(list(dp_semantic_remap_dict.keys()))
dp_remap_dict = denseposelib.semantic_remap_dict2remap_dict(
dp_semantic_remap_dict, dp_new_part_list)
remapped_gt_segmentation, remapped_inferred = denseposelib.get_best_segmentation(
iuvs, labels, dp_remap_dict)
df = pd.DataFrame(columns=["batch_idx"] + dp_new_part_list)
df = denseposelib.calculate_iou_df(remapped_inferred,
remapped_gt_segmentation,
dp_new_part_list)
df.to_csv(os.path.join(output_folder, "part_ious.csv"),
index=False,
header=True)
df_mean = denseposelib.calculate_overall_iou_from_df(df)
with open(os.path.join(output_folder, "mean_part_ios.csv"), "w") as f:
print(
tabulate(df_mean,
headers="keys",
tablefmt="psql",
showindex="never"),
file=f,
)
target_dir = os.path.join(output_folder, "04_compare")
os.makedirs(target_dir, exist_ok=True)
background_color = np.array([1, 1, 1])
colors1 = imageutils.make_colors(config["n_inferred_parts"] + 1,
with_background=True,
background_id=0)
colors2 = imageutils.make_colors(
len(dp_new_part_list),
with_background=True,
background_id=dp_new_part_list.index("background"))
for i, (im1, im2, im3) in enumerate(
zip(labels, remapped_inferred, remapped_gt_segmentation)):
canvas = np.concatenate([colors1[im1], colors2[im2], colors2[im3]],
1).astype(np.float32)
canvas = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR)
fname = os.path.join(target_dir, "{:06d}.png".format(i))
cv2.imwrite(fname, imageutils.convert_range(canvas, [0, 1], [0, 255]))
# batches.plot_batch(
# imageutils.convert_range(canvas, [0, 1], [-1, 1]), fname, cols=3
# )
target_dir = os.path.join(output_folder, "05_remapped_inferred")
os.makedirs(target_dir, exist_ok=True)
write_labels(remapped_inferred, target_dir, colors2, n_processes)
target_dir = os.path.join(output_folder, "06_remapped_labels")
os.makedirs(target_dir, exist_ok=True)
write_labels(remapped_gt_segmentation, target_dir, colors2, n_processes)