File size: 5,991 Bytes
b13b124 |
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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import argparse
import gzip
import os
import os.path as osp
import tarfile
import tempfile
import mmcv
STARE_LEN = 20
TRAINING_LEN = 10
def un_gz(src, dst):
g_file = gzip.GzipFile(src)
with open(dst, 'wb+') as f:
f.write(g_file.read())
g_file.close()
def parse_args():
parser = argparse.ArgumentParser(
description='Convert STARE dataset to mmsegmentation format')
parser.add_argument('image_path', help='the path of stare-images.tar')
parser.add_argument('labels_ah', help='the path of labels-ah.tar')
parser.add_argument('labels_vk', help='the path of labels-vk.tar')
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 main():
args = parse_args()
image_path = args.image_path
labels_ah = args.labels_ah
labels_vk = args.labels_vk
if args.out_dir is None:
out_dir = osp.join('data', 'STARE')
else:
out_dir = args.out_dir
print('Making directories...')
mmcv.mkdir_or_exist(out_dir)
mmcv.mkdir_or_exist(osp.join(out_dir, 'images'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'images', 'training'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'images', 'validation'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations', 'training'))
mmcv.mkdir_or_exist(osp.join(out_dir, 'annotations', 'validation'))
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'gz'))
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'files'))
print('Extracting stare-images.tar...')
with tarfile.open(image_path) as f:
f.extractall(osp.join(tmp_dir, 'gz'))
for filename in os.listdir(osp.join(tmp_dir, 'gz')):
un_gz(
osp.join(tmp_dir, 'gz', filename),
osp.join(tmp_dir, 'files',
osp.splitext(filename)[0]))
now_dir = osp.join(tmp_dir, 'files')
assert len(os.listdir(now_dir)) == STARE_LEN, \
'len(os.listdir(now_dir)) != {}'.format(STARE_LEN)
for filename in sorted(os.listdir(now_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img,
osp.join(out_dir, 'images', 'training',
osp.splitext(filename)[0] + '.png'))
for filename in sorted(os.listdir(now_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img,
osp.join(out_dir, 'images', 'validation',
osp.splitext(filename)[0] + '.png'))
print('Removing the temporary files...')
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'gz'))
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'files'))
print('Extracting labels-ah.tar...')
with tarfile.open(labels_ah) as f:
f.extractall(osp.join(tmp_dir, 'gz'))
for filename in os.listdir(osp.join(tmp_dir, 'gz')):
un_gz(
osp.join(tmp_dir, 'gz', filename),
osp.join(tmp_dir, 'files',
osp.splitext(filename)[0]))
now_dir = osp.join(tmp_dir, 'files')
assert len(os.listdir(now_dir)) == STARE_LEN, \
'len(os.listdir(now_dir)) != {}'.format(STARE_LEN)
for filename in sorted(os.listdir(now_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(now_dir, filename))
# The annotation img should be divided by 128, because some of
# the annotation imgs are not standard. We should set a threshold
# to convert the nonstandard annotation imgs. The value divided by
# 128 equivalent to '1 if value >= 128 else 0'
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'training',
osp.splitext(filename)[0] + '.png'))
for filename in sorted(os.listdir(now_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'validation',
osp.splitext(filename)[0] + '.png'))
print('Removing the temporary files...')
with tempfile.TemporaryDirectory(dir=args.tmp_dir) as tmp_dir:
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'gz'))
mmcv.mkdir_or_exist(osp.join(tmp_dir, 'files'))
print('Extracting labels-vk.tar...')
with tarfile.open(labels_vk) as f:
f.extractall(osp.join(tmp_dir, 'gz'))
for filename in os.listdir(osp.join(tmp_dir, 'gz')):
un_gz(
osp.join(tmp_dir, 'gz', filename),
osp.join(tmp_dir, 'files',
osp.splitext(filename)[0]))
now_dir = osp.join(tmp_dir, 'files')
assert len(os.listdir(now_dir)) == STARE_LEN, \
'len(os.listdir(now_dir)) != {}'.format(STARE_LEN)
for filename in sorted(os.listdir(now_dir))[:TRAINING_LEN]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'training',
osp.splitext(filename)[0] + '.png'))
for filename in sorted(os.listdir(now_dir))[TRAINING_LEN:]:
img = mmcv.imread(osp.join(now_dir, filename))
mmcv.imwrite(
img[:, :, 0] // 128,
osp.join(out_dir, 'annotations', 'validation',
osp.splitext(filename)[0] + '.png'))
print('Removing the temporary files...')
print('Done!')
if __name__ == '__main__':
main()
|