climategan / utils_scripts /merge_labelbox_masks.py
NimaBoscarino's picture
copy the climategan repo in here
6e601ed
raw
history blame
1.44 kB
from pathlib import Path
import numpy as np
from skimage.io import imread, imsave
from shutil import copyfile
if __name__ == "__main__":
# output of download_labelbox.py
base_dir = Path("/Users/victor/Downloads/labelbox_test_flood-v2")
labeled_dir = base_dir / "__labeled"
assert base_dir.exists()
labeled_dir.mkdir(exist_ok=True)
sub_dirs = [
d
for d in base_dir.expanduser().resolve().iterdir()
if d.is_dir() and not d.name.startswith(".") and d.name != "__labeled"
]
for k, sd in enumerate(sub_dirs):
print(k + 1, "/", len(sub_dirs), sd.name)
# must-flood binary mask
must = np.stack([imread(i)[:, :, :3] for i in sd.glob("*must*.png")]).sum(0) > 0
# cannot-flood binary mask
cannot = (
np.stack([imread(i)[:, :, :3] for i in sd.glob("*cannot*.png")]).sum(0) > 0
)
# must is red
must = (must * [0, 0, 255]).astype(np.uint8)
# connot is blue
cannot = (cannot * [255, 0, 0]).astype(np.uint8)
# merged labels
label = must + cannot
# check no overlap
assert sorted(np.unique(label)) == [0, 255]
# create filename
stem = "_".join(list(sd.glob("*must*.png"))[0].stem.split("_")[:-2])
# save label
imsave(sd / f"{stem}_labeled.png", label)
copyfile(sd / f"{stem}_labeled.png", labeled_dir / f"{stem}_labeled.png")