Spaces:
Runtime error
Runtime error
File size: 2,839 Bytes
3b96cb1 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import argparse
import os.path as osp
import shutil
import tempfile
import zipfile
from mmengine.utils import mkdir_or_exist
def parse_args():
parser = argparse.ArgumentParser(
description='Convert NYU Depth dataset to mmsegmentation format')
parser.add_argument('raw_data', help='the path of raw data')
parser.add_argument(
'-o', '--out_dir', help='output path', default='./data/nyu')
args = parser.parse_args()
return args
def reorganize(raw_data_dir: str, out_dir: str):
"""Reorganize NYU Depth dataset files into the required directory
structure.
Args:
raw_data_dir (str): Path to the raw data directory.
out_dir (str): Output directory for the organized dataset.
"""
def move_data(data_list, dst_prefix, fname_func):
"""Move data files from source to destination directory.
Args:
data_list (list): List of data file paths.
dst_prefix (str): Prefix to be added to destination paths.
fname_func (callable): Function to process file names
"""
for data_item in data_list:
data_item = data_item.strip().strip('/')
new_item = fname_func(data_item)
shutil.move(
osp.join(raw_data_dir, data_item),
osp.join(out_dir, dst_prefix, new_item))
def process_phase(phase):
"""Process a dataset phase (e.g., 'train' or 'test')."""
with open(osp.join(raw_data_dir, f'nyu_{phase}.txt')) as f:
data = filter(lambda x: len(x.strip()) > 0, f.readlines())
data = map(lambda x: x.split()[:2], data)
images, annos = zip(*data)
move_data(images, f'images/{phase}',
lambda x: x.replace('/rgb', ''))
move_data(annos, f'annotations/{phase}',
lambda x: x.replace('/sync_depth', ''))
process_phase('train')
process_phase('test')
def main():
args = parse_args()
print('Making directories...')
mkdir_or_exist(args.out_dir)
for subdir in [
'images/train', 'images/test', 'annotations/train',
'annotations/test'
]:
mkdir_or_exist(osp.join(args.out_dir, subdir))
print('Generating images and annotations...')
if args.raw_data.endswith('.zip'):
with tempfile.TemporaryDirectory() as tmp_dir:
zip_file = zipfile.ZipFile(args.raw_data)
zip_file.extractall(tmp_dir)
reorganize(osp.join(tmp_dir, 'nyu'), args.out_dir)
else:
assert osp.isdir(
args.raw_data
), 'the argument --raw-data should be either a zip file or directory.'
reorganize(args.raw_data, args.out_dir)
print('Done!')
if __name__ == '__main__':
main()
|