Some scripts for medical image processing

utils.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Mar 28 14:34:17 2018

@author: qzane
"""

import os,sys
import numpy as np

import torch

from PIL import Image
import SimpleITK as sitk


def itk_nib(mat):
    '''convert between SimpleITK and Nibabel format'''
    mat = mat.squeeze()
    return mat.transpose((2,1,0))

nib_itk = itk_nib

def itk_save(mat, file, template):
    img = sitk.GetImageFromArray(mat)

    img.SetDirection(template.GetDirection())
    img.SetOrigin(template.GetOrigin())
    img.SetSpacing(template.GetSpacing())
    
    sitk.WriteImage(img, file)
    
def itk_read(file):
    im = sitk.ReadImage(file)
    da = sitk.GetArrayFromImage(im).squeeze()
    
    return (im, da)


def bounding_box(mat):
    ''' get the bounding box of an numpy ndarray '''
    index = np.array(np.where(mat>0))
    maxs = index.max(1)
    mins = index.min(1)
    res = np.array(tuple(zip(mins, maxs)), np.int32)
    return res


def hist_match(img, temp):
    ''' histogram matching from img to temp '''
    matcher = sitk.HistogramMatchingImageFilter()
    matcher.SetNumberOfHistogramLevels(1024)
    matcher.SetNumberOfMatchPoints(7)
    matcher.ThresholdAtMeanIntensityOn()
    res = matcher.Execute(img,temp)
    return res

def resize(mat, zoom=1.0, mode='AA', rotate1=0, rotate2=0):
    '''mode == 'AA' /'ANTIALIAS' or 'NN' / 'NEAREST'''
    mat = mat.squeeze()
    
    dtype = mat.dtype
    if mode.upper() in ('AA', 'ANTIALIAS'):
        mode = Image.ANTIALIAS
        mat = np.float64(mat)
    elif mode.upper() in ('NN', 'NEAREST'):
        mode = Image.NEAREST
        mat = np.float64(mat>0)*100
    else:
        raise(Exception('unexpected resize mode'))
        
    shape = mat.shape
        
    if type(zoom) == tuple:
        shape1 = (shape[0],) +zoom[1:]
        shape2 = zoom
    else:
        shape1 = (shape[0], int(shape[1]*zoom), int(shape[2]*zoom))
        shape2 = (int(shape[0]*zoom), shape1[1], shape1[2])

    mat1 = np.zeros(shape1, mat.dtype)
    mat2 = np.zeros(shape2, mat.dtype)
    
    
    for no in range(shape[0]):
        mat1[no] = np.array(Image.fromarray(mat[no]).resize(shape1[:0:-1], mode).rotate(rotate2))
        
    for no in range(shape2[2]):
        mat2[:,:,no] = np.array(Image.fromarray(mat1[:,:,no]).resize(shape2[1::-1],
                                                                     mode).rotate(rotate1))
    if mode == Image.NEAREST:
        mat2 = (mat2 > 0)
    
    return mat2.astype(dtype)

def resize_mask(mat, zoom=1.0,):
    dtype = mat.dtype
    mat = np.float32(mat)
    mat = mat * 100
    mat = resize(mat, zoom, 'AA')
    mat = mat / 100 + 0.5
    mat = mat.astype(dtype)
    return mat
    


def fitMat(targetshape, mat, center=None, target_center=None):
    ''' fit the mat into targetshape
    make sure the center of the origin image locat in the target_center of the new mat '''
    
    
    if not isinstance(mat, np.ndarray):
        raise(Exception('tensor type not supported'))
        
    
    mat = mat.squeeze()
    if center is None:
        center = tuple(map(lambda x:x//2, mat.shape))
    else:
        assert(len(center)==3)
    if target_center is None:
        target_center = tuple(map(lambda x:x//2, targetshape))
    else:
        assert(len(center)==3)
        

    
    c1 = target_center
    c2 = center
    
    outMat = np.zeros(targetshape, mat.dtype)
    outMatSqueezed = outMat.squeeze()
    
    r3 = (min(c1[0],c2[0]),
          min(c1[1],c2[1]),   
          min(c1[2],c2[2]))
    
    r4 = (min(mat.shape[0]-c2[0],targetshape[0]-c1[0]),
          min(mat.shape[1]-c2[1],targetshape[1]-c1[1]),
          min(mat.shape[2]-c2[2],targetshape[2]-c1[2]))
    
    outMatSqueezed[c1[0]-r3[0]:c1[0]+r4[0],
           c1[1]-r3[1]:c1[1]+r4[1],
           c1[2]-r3[2]:c1[2]+r4[2]] = \
        mat[c2[0]-r3[0]:c2[0]+r4[0],
            c2[1]-r3[1]:c2[1]+r4[1],
            c2[2]-r3[2]:c2[2]+r4[2]]
        
    return outMat

def fitMat_test():
    mat1 = np.zeros((321,345,367))
    mat1[11,22,33] = 1
    mat2 = fitMat((431,275,367), mat1, (11,22,33), (430, 270, 11))
    where = np.where(mat2>0)
    if((where[0][0]==430) and (where[1][0]==270) and (where[2][0]==11)):
        return True
    else:
        return False

assert(fitMat_test)

def makeSlice(mat, piece, targetshape, step=32, overlap=32):
    ''' when piece == -1, return all mats
    '''
    
    tmpMat = fitMat(targetshape, mat)
            
    if piece == -1: # return all mats
        res = []
        for x in range((targetshape[0]-overlap) // step):
            for y in range((targetshape[0]-overlap) // step):
                for z in range((targetshape[0]-overlap) // step):
                    xx = x*step
                    yy = y*step
                    zz = z*step
                    res.append(tmpMat[xx:xx+overlap+step,
                                      yy:yy+overlap+step,
                                      zz:zz+overlap+step])
        return res
    
    no = 0
    for x in range((targetshape[0]-overlap) // step):
        for y in range((targetshape[0]-overlap) // step):
            for z in range((targetshape[0]-overlap) // step):
                if no == piece:
                    xx = x*step
                    yy = y*step
                    zz = z*step
                    return tmpMat[xx:xx+overlap+step,
                                  yy:yy+overlap+step,
                                  zz:zz+overlap+step]
                else:
                    no += 1


def combineSlice(mats, targetshape, step=32, overlap=32):
    ''' mats are probability maps '''
    
    outMat = np.zeros(targetshape, mats[0].dtype)
    countMat = np.zeros(targetshape, mats[0].dtype)
    
    no = 0
    for x in range((targetshape[0]-overlap) // step):
        for y in range((targetshape[0]-overlap) // step):
            for z in range((targetshape[0]-overlap) // step):
                xx = x*step
                yy = y*step
                zz = z*step
                countMat[xx:xx+overlap+step,
	                 yy:yy+overlap+step,
	                 zz:zz+overlap+step] += 1
                outMat[xx:xx+overlap+step,
	               yy:yy+overlap+step,
	               zz:zz+overlap+step] += mats[no]
                no += 1
        
    countMat[countMat==0]=1
    outMat = outMat / countMat
    return outMat


def mat2tensor(im_resize, std=None):
    im_resize = im_resize.squeeze()
    im_resize = np.float32(im_resize)
    if std is None:
       std = np.where(np.isclose(im_resize,0), np.nan, im_resize)
    if std == 0:
        pass #im_resize = im_resize / 1
    else:
        im_resize = im_resize / std
        
    #im_resize = im_resize - im_resize.mean() ##really???
    #im_resize = im_resize / im_resize.std() 
    
    im_resize = torch.from_numpy(im_resize)
    im_resize = im_resize.contiguous().unsqueeze(0)
    return im_resize



def change_spacing(mat, old_spacing, new_spacing, itk_spacing=True):
    ''' resize mat so the old spacing changes to the new spacing '''
    if itk_spacing:
        mat_trans = mat.transpose((2,1,0)) # the order of mat and spacing is different
        
    shape = (int(mat_trans.shape[0]*old_spacing[0]/new_spacing[0]),
             int(mat_trans.shape[1]*old_spacing[1]/new_spacing[1]),
             int(mat_trans.shape[2]*old_spacing[2]/new_spacing[2]))
    new_mat = resize(mat_trans, shape)
    
    if itk_spacing:
        new_mat = new_mat.transpose((2,1,0)) # the order of mat and spacing is different
    return new_mat
             
def nonzero_mean_std(mat):
    ''' calculate mean and std for non-zero pixels only '''
    tmp = np.where(np.isclose(mat,0), np.nan, mat)
    return np.nanmean(tmp), np.nanstd(tmp)