Spaces:
Runtime error
Runtime error
File size: 3,178 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 |
# 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()
|