Spaces:
Running
Running
| # ------------------------------------------------------------------------------ | |
| # 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") |