File size: 2,420 Bytes
d617811 |
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 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
# Copyright (c) Facebook, Inc. and its affiliates.
# Copyright (c) Meta Platforms, Inc. All Rights Reserved
# Modified by Feng Liang from https://github.com/MendelXu/zsseg.baseline/blob/master/datasets/prepare_voc_sem_seg.py
# Modified by Heeseong Shin from https://github.com/facebookresearch/ov-seg/blob/main/datasets/prepare_voc_sem_seg.py
import os
import os.path as osp
from pathlib import Path
import tqdm
import numpy as np
from PIL import Image
clsID_to_trID = {
0: 255,
1: 0,
2: 1,
3: 2,
4: 3,
5: 4,
6: 5,
7: 6,
8: 7,
9: 8,
10: 9,
11: 10,
12: 11,
13: 12,
14: 13,
15: 14,
16: 15,
17: 16,
18: 17,
19: 18,
20: 19,
255: 255,
}
clsID_to_trID_bg = clsID_to_trID.copy()
clsID_to_trID_bg[0] = 20
def convert_to_trainID(
maskpath, out_mask_dir, is_train, clsID_to_trID=clsID_to_trID, suffix=""
):
mask = np.array(Image.open(maskpath))
mask_copy = np.ones_like(mask, dtype=np.uint8) * 255
for clsID, trID in clsID_to_trID.items():
mask_copy[mask == clsID] = trID
seg_filename = (
osp.join(out_mask_dir, "train" + suffix, osp.basename(maskpath))
if is_train
else osp.join(out_mask_dir, "val" + suffix, osp.basename(maskpath))
)
if len(np.unique(mask_copy)) == 1 and np.unique(mask_copy)[0] == 255:
return
Image.fromarray(mask_copy).save(seg_filename, "PNG")
if __name__ == "__main__":
dataset_dir = Path(os.getenv("DETECTRON2_DATASETS", "datasets"))
print('Caution: we only generate the validation set!')
voc_path = dataset_dir / "VOCdevkit" / "VOC2012"
out_mask_dir = voc_path / "annotations_detectron2"
out_mask_dir_bg = voc_path / "annotations_detectron2_bg"
#out_image_dir = voc_path / "images_detectron2"
for name in ["val"]:
os.makedirs((out_mask_dir / name), exist_ok=True)
os.makedirs((out_mask_dir_bg / name), exist_ok=True)
#os.makedirs((out_image_dir / name), exist_ok=True)
val_list = [
osp.join(voc_path, "SegmentationClassAug", f + ".png")
for f in np.loadtxt(osp.join(voc_path, "ImageSets/Segmentation/val.txt"), dtype=np.str).tolist()
]
for file in tqdm.tqdm(val_list):
convert_to_trainID(file, out_mask_dir, is_train=False)
convert_to_trainID(file, out_mask_dir_bg, is_train=False, clsID_to_trID=clsID_to_trID_bg) |