# tools/dataset_converters/gid.py
import argparse
import glob
import math
import os
import os.path as osp
from PIL import Image
import mmcv
import numpy as np
from mmengine.utils import ProgressBar, mkdir_or_exist
def parse_args():
parser = argparse.ArgumentParser(
description='Convert GID dataset to mmsegmentation format')
parser.add_argument('dataset_img_path', help='GID images folder path')
parser.add_argument('dataset_label_path', help='GID labels folder path')
parser.add_argument('--tmp_dir', help='path of the temporary directory')
parser.add_argument('-o', '--out_dir', help='output path', default='data/gid')
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
GID_COLORMAP = dict(
Background=(0, 0, 0), #0-背景-黑色
Building=(255, 0, 0), #1-建筑-红色
Farmland=(0, 255, 0), #2-农田-绿色
Forest=(0, 0, 255), #3-森林-蓝色
Meadow=(255, 255, 0),#4-草地-黄色
Water=(0, 0, 255)#5-水-蓝色
)
palette = list(GID_COLORMAP.values())
classes = list(GID_COLORMAP.keys())
#############用列表来存一个 RGB 和一个类别的对应################
def colormap2label(palette):
colormap2label_list = np.zeros(256**3, dtype = np.longlong)
for i, colormap in enumerate(palette):
colormap2label_list[(colormap[0] * 256 + colormap[1])*256+colormap[2]] = i
return colormap2label_list
#############给定那个列表，和vis_png然后生成masks_png################
def label_indices(RGB_label, colormap2label_list):
RGB_label = RGB_label.astype('int32')
idx = (RGB_label[:, :, 0] * 256 + RGB_label[:, :, 1]) * 256 + RGB_label[:, :, 2]
# print(idx.shape)
return colormap2label_list[idx]
def RGB2mask(RGB_label, colormap2label_list):
# RGB_label = np.array(Image.open(RGB_label).convert('RGB')) #打开RGB_png
mask_label = label_indices(RGB_label, colormap2label_list) # .numpy()
return mask_label
colormap2label_list = colormap2label(palette)
def clip_big_image(image_path, clip_save_dir, args, to_label=False):
"""
Original image of GID dataset is very large, thus pre-processing
of them is adopted. Given fixed clip size and stride size to generate
clipped image, the intersection of width and height is determined.
For example, given one 6800 x 7200 original image, the clip size is
256 and stride size is 256, thus it would generate 29 x 27 = 783 images
whose size are all 256 x 256.
"""
image = mmcv.imread(image_path, channel_order='rgb')
# image = mmcv.bgr2gray(image)
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 = RGB2mask(image, colormap2label_list) #这里得改一下
for count, box in enumerate(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, :]
img_name = osp.basename(image_path).replace('.tif', '')
img_name = img_name.replace('_label', '')
if count % 3 == 0:
mmcv.imwrite(
clipped_image.astype(np.uint8),
osp.join(
clip_save_dir.replace('train', 'val'),
f'{img_name}_{start_x}_{start_y}_{end_x}_{end_y}.png'))
else:
mmcv.imwrite(
clipped_image.astype(np.uint8),
osp.join(
clip_save_dir,
f'{img_name}_{start_x}_{start_y}_{end_x}_{end_y}.png'))
count += 1
def main():
args = parse_args()
"""
According to this paper: https://ieeexplore.ieee.org/document/9343296/
select 15 images contained in GID, , which cover the whole six
categories, to generate train set and validation set.
According to Paper: https://ieeexplore.ieee.org/document/9343296/
"""
if args.out_dir is None:
out_dir = osp.join('data', 'gid')
else:
out_dir = args.out_dir
print('Making directories...')
mkdir_or_exist(osp.join(out_dir, 'img_dir', 'train'))
mkdir_or_exist(osp.join(out_dir, 'img_dir', 'val'))
mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'train'))
mkdir_or_exist(osp.join(out_dir, 'ann_dir', 'val'))
src_path_list = glob.glob(os.path.join(args.dataset_img_path, '*.tif'))
print(f'Find {len(src_path_list)} pictures')
prog_bar = ProgressBar(len(src_path_list))
dst_img_dir = osp.join(out_dir, 'img_dir', 'train')
dst_label_dir = osp.join(out_dir, 'ann_dir', 'train')
for i, img_path in enumerate(src_path_list):
label_path = osp.join(args.dataset_label_path, osp.basename(img_path.replace('.tif', '_label.tif')))
clip_big_image(img_path, dst_img_dir, args, to_label=False)
clip_big_image(label_path, dst_label_dir, args, to_label=True)
prog_bar.update()
print('Done!')
if __name__ == '__main__':
main()
