FC-CLIP / datasets /prepare_pascal_voc_sem_seg.py
yucornetto's picture
init for demo
b6396ac
raw
history blame
2.43 kB
# ------------------------------------------------------------------------------
# Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# This work is made available under the Nvidia Source Code License.
# To view a copy of this license, visit
# https://github.com/NVlabs/ODISE/blob/main/LICENSE
#
# Written by Jiarui Xu
# ------------------------------------------------------------------------------
import os
from pathlib import Path
import shutil
import numpy as np
import tqdm
from PIL import Image
def convert_pas21(input, output):
img = np.asarray(Image.open(input))
assert img.dtype == np.uint8
# do nothing
Image.fromarray(img).save(output)
def convert_pas20(input, output):
img = np.array(Image.open(input))
img[img == 0] = 255
img = img - 1
img[img == 254] = 255
assert img.dtype == np.uint8
# do nothing
Image.fromarray(img).save(output)
if __name__ == "__main__":
dataset_dir = Path(os.getenv("DETECTRON2_DATASETS", "datasets")) / "pascal_voc_d2"
voc_dir = Path(os.getenv("DETECTRON2_DATASETS", "datasets")) / "VOCdevkit/VOC2012"
for split in ["training", "validation"]:
if split == "training":
img_name_path = voc_dir / "ImageSets/Segmentation/train.txt"
else:
img_name_path = voc_dir / "ImageSets/Segmentation/val.txt"
img_dir = voc_dir / "JPEGImages"
ann_dir = voc_dir / "SegmentationClass"
output_img_dir = dataset_dir / "images" / split
output_ann_dir_21 = dataset_dir / "annotations_pascal21" / split
output_ann_dir_20 = dataset_dir / "annotations_pascal20" / split
output_img_dir.mkdir(parents=True, exist_ok=True)
output_ann_dir_21.mkdir(parents=True, exist_ok=True)
output_ann_dir_20.mkdir(parents=True, exist_ok=True)
with open(img_name_path) as f:
for line in tqdm.tqdm(f.readlines()):
img_name = line.strip()
img_path = img_dir / f"{img_name}.jpg"
ann_path = ann_dir / f"{img_name}.png"
# print(f'copy2 {output_img_dir}')
shutil.copy2(img_path, output_img_dir)
# print(f"convert {ann_dir} to {output_ann_dir / f'{img_name}.png'}")
convert_pas21(ann_path, output_ann_dir_21 / f"{img_name}.png")
convert_pas20(ann_path, output_ann_dir_20 / f"{img_name}.png")