Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import argparse | |
| import glob | |
| import math | |
| import os | |
| import os.path as osp | |
| import mmcv | |
| import numpy as np | |
| from mmengine.utils import ProgressBar | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='Convert levir-cd dataset to mmsegmentation format') | |
| parser.add_argument('--dataset_path', help='potsdam folder path') | |
| parser.add_argument('-o', '--out_dir', help='output path') | |
| parser.add_argument( | |
| '--clip_size', | |
| type=int, | |
| help='clipped size of image after preparation', | |
| default=256) | |
| parser.add_argument( | |
| '--stride_size', | |
| type=int, | |
| help='stride of clipping original images', | |
| default=256) | |
| args = parser.parse_args() | |
| return args | |
| def main(): | |
| args = parse_args() | |
| input_folder = args.dataset_path | |
| png_files = glob.glob( | |
| os.path.join(input_folder, '**/*.png'), recursive=True) | |
| output_folder = args.out_dir | |
| prog_bar = ProgressBar(len(png_files)) | |
| for png_file in png_files: | |
| new_path = os.path.join( | |
| output_folder, | |
| os.path.relpath(os.path.dirname(png_file), input_folder)) | |
| os.makedirs(os.path.dirname(new_path), exist_ok=True) | |
| label = False | |
| if 'label' in png_file: | |
| label = True | |
| clip_big_image(png_file, new_path, args, label) | |
| prog_bar.update() | |
| def clip_big_image(image_path, clip_save_dir, args, to_label=False): | |
| image = mmcv.imread(image_path) | |
| h, w, c = image.shape | |
| clip_size = args.clip_size | |
| stride_size = args.stride_size | |
| num_rows = math.ceil((h - clip_size) / stride_size) if math.ceil( | |
| (h - clip_size) / | |
| stride_size) * stride_size + clip_size >= h else math.ceil( | |
| (h - clip_size) / stride_size) + 1 | |
| num_cols = math.ceil((w - clip_size) / stride_size) if math.ceil( | |
| (w - clip_size) / | |
| stride_size) * stride_size + clip_size >= w else math.ceil( | |
| (w - clip_size) / stride_size) + 1 | |
| x, y = np.meshgrid(np.arange(num_cols + 1), np.arange(num_rows + 1)) | |
| xmin = x * clip_size | |
| ymin = y * clip_size | |
| xmin = xmin.ravel() | |
| ymin = ymin.ravel() | |
| xmin_offset = np.where(xmin + clip_size > w, w - xmin - clip_size, | |
| np.zeros_like(xmin)) | |
| ymin_offset = np.where(ymin + clip_size > h, h - ymin - clip_size, | |
| np.zeros_like(ymin)) | |
| boxes = np.stack([ | |
| xmin + xmin_offset, ymin + ymin_offset, | |
| np.minimum(xmin + clip_size, w), | |
| np.minimum(ymin + clip_size, h) | |
| ], | |
| axis=1) | |
| if to_label: | |
| image[image == 255] = 1 | |
| image = image[:, :, 0] | |
| for box in boxes: | |
| start_x, start_y, end_x, end_y = box | |
| clipped_image = image[start_y:end_y, start_x:end_x] \ | |
| if to_label else image[start_y:end_y, start_x:end_x, :] | |
| idx = osp.basename(image_path).split('.')[0] | |
| mmcv.imwrite( | |
| clipped_image.astype(np.uint8), | |
| osp.join(clip_save_dir, | |
| f'{idx}_{start_x}_{start_y}_{end_x}_{end_y}.png')) | |
| if __name__ == '__main__': | |
| main() | |