Spaces:
Runtime error
Runtime error
File size: 3,773 Bytes
412c852 |
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 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os
import os.path as osp
import tempfile
import zipfile
import mmcv
import numpy as np
from mmengine.utils import mkdir_or_exist
def parse_args():
parser = argparse.ArgumentParser(
description='Convert REFUGE dataset to mmsegmentation format')
parser.add_argument('--raw_data_root', help='the root path of raw data')
parser.add_argument('--tmp_dir', help='path of the temporary directory')
parser.add_argument('-o', '--out_dir', help='output path')
args = parser.parse_args()
return args
def extract_img(root: str,
cur_dir: str,
out_dir: str,
mode: str = 'train',
file_type: str = 'img') -> None:
"""_summary_
Args:
Args:
root (str): root where the extracted data is saved
cur_dir (cur_dir): dir where the zip_file exists
out_dir (str): root dir where the data is saved
mode (str, optional): Defaults to 'train'.
file_type (str, optional): Defaults to 'img',else to 'mask'.
"""
zip_file = zipfile.ZipFile(cur_dir)
zip_file.extractall(root)
for cur_dir, dirs, files in os.walk(root):
# filter child dirs and directories with "Illustration" and "MACOSX"
if len(dirs) == 0 and \
cur_dir.split('\\')[-1].find('Illustration') == -1 and \
cur_dir.find('MACOSX') == -1:
file_names = [
file for file in files
if file.endswith('.jpg') or file.endswith('.bmp')
]
for filename in sorted(file_names):
img = mmcv.imread(osp.join(cur_dir, filename))
if file_type == 'annotations':
img = img[:, :, 0]
img[np.where(img == 0)] = 1
img[np.where(img == 128)] = 2
img[np.where(img == 255)] = 0
mmcv.imwrite(
img,
osp.join(out_dir, file_type, mode,
osp.splitext(filename)[0] + '.png'))
def main():
args = parse_args()
raw_data_root = args.raw_data_root
if args.out_dir is None:
out_dir = osp.join('./data', 'REFUGE')
else:
out_dir = args.out_dir
print('Making directories...')
mkdir_or_exist(out_dir)
mkdir_or_exist(osp.join(out_dir, 'images'))
mkdir_or_exist(osp.join(out_dir, 'images', 'training'))
mkdir_or_exist(osp.join(out_dir, 'images', 'validation'))
mkdir_or_exist(osp.join(out_dir, 'images', 'test'))
mkdir_or_exist(osp.join(out_dir, 'annotations'))
mkdir_or_exist(osp.join(out_dir, 'annotations', 'training'))
mkdir_or_exist(osp.join(out_dir, 'annotations', 'validation'))
mkdir_or_exist(osp.join(out_dir, 'annotations', 'test'))
print('Generating images and annotations...')
# process data from the child dir on the first rank
cur_dir, dirs, files = list(os.walk(raw_data_root))[0]
print('====================')
files = list(filter(lambda x: x.endswith('.zip'), files))
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
for file in files:
# search data folders for training,validation,test
mode = list(
filter(lambda x: file.lower().find(x) != -1,
['training', 'test', 'validation']))[0]
file_root = osp.join(tmp_dir, file[:-4])
file_type = 'images' if file.find('Anno') == -1 and file.find(
'GT') == -1 else 'annotations'
extract_img(file_root, osp.join(cur_dir, file), out_dir, mode,
file_type)
print('Done!')
if __name__ == '__main__':
main()
|