Three-dimensional medical image--two three-dimensional nii images combined into a checkerboard Python

itkwidgets.checkboard or vtkRectilinearWipWidget both refer to two-dimensional images. (There may also be ones for 3D, but I couldn't find them).
So, I wrote my own that does a checkerboard format combination of two nii (3D) images .
Disadvantages : When the image size and spacing are not proportional, the excess part is not within the checkerboard.

Use plt to display the picture first, check the effect , and save the nii picture if you are satisfied.

import torch
import torch.nn.functional as F
import os
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def save_img_nii(I_img, savename,affine=None,header=None):
    # 1.5代表图像voxel Spacing（此处自行修改）
    affine = np.diag([1.5, 1.5, 1.5, 1])
    new_img = nib.nifti1.Nifti1Image(I_img, affine, header)
    nib.save(new_img, savename)

def get_img_nii(img_path,inshape=None):
    img_org = nib.load(img_path).get_fdata()
    org_x, org_y, org_z = img_org.shape[-3], img_org.shape[-2], img_org.shape[-1]
    img_size=[org_x, org_y, org_z]

    # resize图像
    if inshape is not None:
        img = torch.tensor(img_org).unsqueeze(0).unsqueeze(0).to(torch.float32)
        img = torch.nn.functional.interpolate(img, size=inshape, mode='trilinear', align_corners=False)
        return img[0,0,...].detach().numpy(),img_size
    else:
        return img_org, img_size

def get_mask_nii(img_path,inshape=None):
    img_org = nib.load(img_path).get_fdata()
    org_x, org_y, org_z = img_org.shape[-3], img_org.shape[-2], img_org.shape[-1]
    img_size=[org_x, org_y, org_z]

    # resize图像
    if inshape is not None:
        img = torch.tensor(img_org).unsqueeze(0).unsqueeze(0).to(torch.float32)
        img = torch.nn.functional.interpolate(img, size=inshape, mode='nearest')
        return img[0,0,...].detach().numpy(),img_size
    else:
        return img_org, img_size

def get_checkboard(size,ker):
    # 先拼出一个二维棋盘，再拼出一个三维棋盘
    checkboard_img2d=np.zeros((size[0],size[1]))
    for i in range(size[0]// ker):  # 计算x轴有几个白块
        for j in range(size[1] // ker):  # 计算y轴有几个白块
                if (i + j) % 2 == 0:
                    row_begin = i * ker
                    row_end = row_begin + ker
                    col_begin = j * ker
                    col_end = col_begin + ker
                    checkboard_img2d[row_begin:row_end, col_begin:col_end] = 1  # 画白块

    checkboard_img2d_inv=abs(checkboard_img2d-1)
    checkboard_img3d = np.zeros((size[0], size[1], size[2]))
    for z in range(size[2]//ker):
        if z%2 == 0:
            for n in range(ker):
                checkboard_img3d[:, :, z*ker+n] = checkboard_img2d
        else:
            for n in range(ker):
                checkboard_img3d[:, :, z*ker+n] = checkboard_img2d_inv

    return checkboard_img3d,abs(checkboard_img3d-1)

def checkboard(img1,img2,ker,size):
    checkboard_3d,checkboard_3d_inv=get_checkboard(size,ker)
    img1_checkboard=img1*checkboard_3d
    img2_checkboard=img2*checkboard_3d_inv
    merge=img1_checkboard+img2_checkboard
    return merge

def show_img(img):
    '''
	显示各维中间一层作为输出。
	'''
    x, y, z = img.shape[-3], fix.shape[-2], fix.shape[-1]
    plt.figure()
    plt.subplot(1, 3, 1)
    img1=np.rot90 (img[x//2,:,:], k = 1, axes = (0, 1)) # 我的图片需要旋转一下看更舒服,可以省略
    plt.imshow(img1,cmap=plt.get_cmap('gray'))
    plt.subplot(1, 3, 2)
    img2=np.rot90 (img[:,y//2,:], k = 1, axes = (0, 1))
    plt.imshow(img2,cmap=plt.get_cmap('gray'))
    plt.subplot(1, 3, 3)
    img3 = np.rot90(img[:,:,z//2], k=1, axes=(0, 1))
    plt.imshow(img3,cmap=plt.get_cmap('gray'))
    plt.show()

if __name__ == '__main__':

    # 获取三维图片
    fix_path='CBCT.nii.gz'
    mov_path='CT.nii.gz'
    fix ,fix_size= get_img_nii(fix_path)
    mov ,mov_size= get_img_nii(mov_path)
    assert fix_size == mov_size, ' two images must have same size'

    # 输入两图形成棋盘图
    ker=8 # 间隔大小
    checkboard_img=checkboard(fix, mov, ker,fix_size)
    show_img(checkboard_img)# 先看一下满不满意，满足则保存
    checkboard_img_name='checkboard_img.nii.gz'
    save_img_nii(checkboard_img, checkboard_img_name)

Note : When the input image is not a multiple of the interval ker, the extra parts of the x and y axes are not spaced. Suggested readers can optimize it by themselves.

plt display checkerboard style

nii save map