语义分割data augmentation/ 三维数据集

• 由于各种包（pytorch等）支持的大多数数据增广是二维的图像，缺少三维图像data/label一起进行变换的增广接口，目前在做生物图像nii.gz格式的数据增广（z方向对齐），使data与label一起旋转任意随机度数。目前写了旋转接口，后续可以增加别的增广方式。

旋转

# -*- coding: utf-8 -*-
# @Time    : 2019/8/10 21:36
# @Author  : Mingxing Li
# @FileName: augmentation.py
# @Software: PyCharm

import numpy as np
import nibabel as nib
import cv2
import os


out_path = r"E:\structseg2019\dataset\StructSeg 2019\Task2_Naso_GTV\augmentation"
file_name = ["anchor_data.nii.gz",
             "anchor_label.nii.gz"]

def load(path):
    data = nib.load(path).get_data()
    print(data.shape)
    return data


data = load(r"E:\structseg2019\dataset\StructSeg 2019\Task2_Naso_GTV\new_crop\81\anchor_data.nii.gz")
label = load(r"E:\structseg2019\dataset\StructSeg 2019\Task2_Naso_GTV\new_crop\81\anchor_label.nii.gz")


def rotate(image, label, angle, center=None, scale=1.0):
    # 获取图像尺寸
    (h, w) = image.shape[:2]

    # 若未指定旋转中心，则将图像中心设为旋转中心
    if center is None:
        center = (w / 2, h / 2)

    # 执行旋转
    M = cv2.getRotationMatrix2D(center, angle, scale)
    rotated_image = cv2.warpAffine(image, M, (w, h))
    rotated_label = cv2.warpAffine(label, M, (w, h))
    # 返回旋转后的图像
    return rotated_image, rotated_label


def nii_gz_operate(data, label):
    angle = np.random.uniform(-7, 7)
    rotated_image, rotated_label = rotate(data, label, angle)
    return rotated_image, rotated_label


rotate_data, rotate_label = nii_gz_operate(data, label)


def save_nii_gz(data, label, fold_seq):
    # w, h, z
    data = data.astype(np.int16)
    data = nib.Nifti1Image(data, np.eye(4))
    label = label.astype(np.int16)
    label = nib.Nifti1Image(label, np.eye(4))

    if str(fold_seq) not in os.listdir(out_path):
        os.mkdir(out_path+"/"+str(fold_seq))

    nib.save(data, out_path+"/"+str(fold_seq)+"/"+file_name[0])
    print(out_path+"/"+str(fold_seq)+"/"+file_name[0], "has loaded! ")

    nib.save(label, out_path+"/"+str(fold_seq)+"/"+file_name[1])
    print(out_path+"/"+str(fold_seq)+"/"+file_name[1], "has loaded! ")


save_nii_gz(rotate_data, rotate_label, 1)