遥感语义分割数据集中的切图策略

遥感语义分割数据集中的切图策略

# 切图脚本
import argparse
import glob
import math
import os
import os.path as osp
import tempfile
import zipfile

import mmcv
import numpy as np
from mmengine.utils import ProgressBar, mkdir_or_exist

def clip_big_image(image_path, clip_save_dir, to_label=False):
    # Original image of Vaihingen 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 5120 x 5120 original image, the clip size is
    # 512 and stride size is 256, thus it would generate 20x20 = 400 images
    # whose size are all 512x512.
    image = mmcv.imread(image_path)

    h, w, c = image.shape
    cs = 512 # todo 分割的大小
    ss = 256 # todo 分割的步长

    num_rows = math.ceil((h - cs) / ss) if math.ceil(
        (h - cs) / ss) * ss + cs >= h else math.ceil((h - cs) / ss) + 1
    num_cols = math.ceil((w - cs) / ss) if math.ceil(
        (w - cs) / ss) * ss + cs >= w else math.ceil((w - cs) / ss) + 1

    x, y = np.meshgrid(np.arange(num_cols + 1), np.arange(num_rows + 1))
    xmin = x * cs
    ymin = y * cs

    xmin = xmin.ravel()
    ymin = ymin.ravel()
    xmin_offset = np.where(xmin + cs > w, w - xmin - cs, np.zeros_like(xmin))
    ymin_offset = np.where(ymin + cs > h, h - ymin - cs, np.zeros_like(ymin))
    boxes = np.stack([
        xmin + xmin_offset, ymin + ymin_offset,
        np.minimum(xmin + cs, w),
        np.minimum(ymin + cs, h)
    ],
                     axis=1)

    # to_label是用来转化label使用的
    if to_label:
        color_map = np.array([[0, 0, 0], [85, 85, 85]])
        flatten_v = np.matmul(
            image.reshape(-1, c),
            np.array([2, 3, 4]).reshape(3, 1))
        out = np.zeros_like(flatten_v)
        for idx, class_color in enumerate(color_map):
            value_idx = np.matmul(class_color,
                                  np.array([2, 3, 4]).reshape(3, 1))
            out[flatten_v == value_idx] = idx
        image = out.reshape(h, w)

    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, :]
        # area_idx = osp.basename(image_path).split('_')[3].strip('.tif')
        # area_idx = osp.basename(image_path).strip('.')
        # print(area_idx)
        # area_idx = osp.basename(image_path).strip('.png')
        mmcv.imwrite(
            clipped_image.astype(np.uint8),
            osp.join(clip_save_dir,
                     # f'{area_idx}_{start_x}_{start_y}_{end_x}_{end_y}.png'))
                     f'img1_{start_x}_{start_y}_{end_x}_{end_y}.png'))

if __name__ == '__main__':
    # 切图搞定，还需要标签转换，和数据集分离
    clip_big_image(image_path="E:/EEEE-COM/toUser/toUser/train/img1/img1.tif", clip_save_dir="E:/EEEE-COM/toUser/toUser/train/split_data/training_images")
    # clip_big_image(image_path="E:/EEEE-COM/toUser/toUser/train/img2/img2.tif", clip_save_dir="E:/EEEE-COM/toUser/toUser/train/split_data/training_images")
    # clip_big_image(image_path="E:/EEEE-COM/toUser/toUser/train/train_labels_png/img1.png", clip_save_dir="E:/EEEE-COM/toUser/toUser/train/split_data/training_labels")
    # clip_big_image(image_path="E:/EEEE-COM/toUser/toUser/train/img2/img2.tif", clip_save_dir="E:/EEEE-COM/toUser/toUser/train/split_data/training_images")