Spaces:
Runtime error
Runtime error
File size: 1,435 Bytes
ce190ee |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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")
|