import torch
import nibabel
import numpy as np
import os, sys, time
import scipy.ndimage
import torch.nn as nn
import torch.nn.functional as F
from numpy.linalg import inv
try:
import resource
except:
pass
# monkey-patch for back-compatibility with older (~1.0.0) torch
try:
import inspect
if not "align_corners" in inspect.signature(F.grid_sample).parameters:
old_grid_sample = torch.nn.functional.grid_sample
F.grid_sample = lambda *x, **k : old_grid_sample(*x)
except:
pass
if len(sys.argv[1:]) == 0:
print("Need to pass one or more T1 image filename as argument")
sys.exit(1)
print("Using all available CPU threads")
if 0: # otherwise, set a limit (useful for running multiple instances)
torch.set_num_threads(4)
class HeadModel(nn.Module):
def __init__(self):
super(HeadModel, self).__init__()
self.conv0a = nn.Conv3d(1, 8, 3, padding=1)
self.conv0b = nn.Conv3d(8, 8, 3, padding=1)
self.bn0a = nn.BatchNorm3d(8)
self.ma1 = nn.MaxPool3d(2)
self.conv1a = nn.Conv3d(8, 16, 3, padding=1)
self.conv1b = nn.Conv3d(16, 24, 3, padding=1)
self.bn1a = nn.BatchNorm3d(24)
self.ma2 = nn.MaxPool3d(2)
self.conv2a = nn.Conv3d(24, 24, 3, padding=1)
self.conv2b = nn.Conv3d(24, 32, 3, padding=1)
self.bn2a = nn.BatchNorm3d(32)
self.ma3 = nn.MaxPool3d(2)
self.conv3a = nn.Conv3d(32, 48, 3, padding=1)
self.conv3b = nn.Conv3d(48, 48, 3, padding=1)
self.bn3a = nn.BatchNorm3d(48)
self.conv2u = nn.Conv3d(48, 24, 3, padding=1)
self.conv2v = nn.Conv3d(24+32, 24, 3, padding=1)
self.bn2u = nn.BatchNorm3d(24)
self.conv1u = nn.Conv3d(24, 24, 3, padding=1)
self.conv1v = nn.Conv3d(24+24, 24, 3, padding=1)
self.bn1u = nn.BatchNorm3d(24)
self.conv0u = nn.Conv3d(24, 16, 3, padding=1)
self.conv0v = nn.Conv3d(16+8, 8, 3, padding=1)
self.bn0u = nn.BatchNorm3d(8)
self.conv1x = nn.Conv3d(8, 4, 1, padding=0)
def forward(self, x):
x = F.elu(self.conv0a(x))
self.li0 = x = F.elu(self.bn0a(self.conv0b(x)))
x = self.ma1(x)
x = F.elu(self.conv1a(x))
self.li1 = x = F.elu(self.bn1a(self.conv1b(x)))
x = self.ma2(x)
x = F.elu(self.conv2a(x))
self.li2 = x = F.elu(self.bn2a(self.conv2b(x)))
x = self.ma3(x)
x = F.elu(self.conv3a(x))
self.li3 = x = F.elu(self.bn3a(self.conv3b(x)))
x = F.interpolate(x, scale_factor=2, mode="nearest")
x = F.elu(self.conv2u(x))
x = torch.cat([x, self.li2], 1)
x = F.elu(self.bn2u(self.conv2v(x)))
self.lo1 = x
x = F.interpolate(x, scale_factor=2, mode="nearest")
x = F.elu(self.conv1u(x))
x = torch.cat([x, self.li1], 1)
x = F.elu(self.bn1u(self.conv1v(x)))
x = F.interpolate(x, scale_factor=2, mode="nearest")
self.la1 = x
x = F.elu(self.conv0u(x))
x = torch.cat([x, self.li0], 1)
x = F.elu(self.bn0u(self.conv0v(x)))
self.out = x = self.conv1x(x)
x = torch.sigmoid(x)
return x
class ModelAff(nn.Module):
def __init__(self):
super(ModelAff, self).__init__()
self.convaff1 = nn.Conv3d(2, 16, 3, padding=1)
self.maaff1 = nn.MaxPool3d(2)
self.convaff2 = nn.Conv3d(16, 16, 3, padding=1)
self.bnaff2 = nn.LayerNorm([32, 32, 32])
self.maaff2 = nn.MaxPool3d(2)
self.convaff3 = nn.Conv3d(16, 32, 3, padding=1)
self.bnaff3 = nn.LayerNorm([16, 16, 16])
self.maaff3 = nn.MaxPool3d(2)
self.convaff4 = nn.Conv3d(32, 64, 3, padding=1)
self.maaff4 = nn.MaxPool3d(2)
self.bnaff4 = nn.LayerNorm([8, 8, 8])
self.convaff5 = nn.Conv3d(64, 128, 1, padding=0)
self.convaff6 = nn.Conv3d(128, 12, 4, padding=0)
gsx, gsy, gsz = 64, 64, 64
gx, gy, gz = np.linspace(-1, 1, gsx), np.linspace(-1, 1, gsy), np.linspace(-1,1, gsz)
grid = np.meshgrid(gx, gy, gz) # Y, X, Z
grid = np.stack([grid[2], grid[1], grid[0], np.ones_like(grid[0])], axis=3)
netgrid = np.swapaxes(grid, 0, 1)[...,[2,1,0,3]]
self.register_buffer('grid', torch.tensor(netgrid.astype("float32"), requires_grad = False))
self.register_buffer('diagA', torch.eye(4, dtype=torch.float32))
def forward(self, outc1):
x = outc1
x = F.relu(self.convaff1(x))
x = self.maaff1(x)
x = F.relu(self.bnaff2(self.convaff2(x)))
x = self.maaff2(x)
x = F.relu(self.bnaff3(self.convaff3(x)))
x = self.maaff3(x)
x = F.relu(self.bnaff4(self.convaff4(x)))
x = self.maaff4(x)
x = F.relu(self.convaff5(x))
x = self.convaff6(x)
x = x.view(-1, 3, 4)
x = torch.cat([x, x[:,0:1] * 0], dim=1)
self.tA = torch.transpose(x + self.diagA, 1, 2)
wgrid = self.grid @ self.tA[:,None,None]
gout = F.grid_sample(outc1, wgrid[...,[2,1,0]], align_corners=True)
return gout, self.tA
def resample_other(self, other):
with torch.no_grad():
wgrid = self.grid @ self.tA[:,None,None]
gout = F.grid_sample(other, wgrid[...,[2,1,0]], align_corners=True)
return gout
def bbox_world(affine, shape):
s = shape[0]-1, shape[1]-1, shape[2]-1
bbox = [[0,0,0], [s[0],0,0], [0,s[1],0], [0,0,s[2]], [s[0],s[1],0], [s[0],0,s[2]], [0,s[1],s[2]], [s[0],s[1],s[2]]]
w = affine @ np.column_stack([bbox, [1]*8]).T
return w.T
bbox_one = np.array([[-1,-1,-1,1], [1, -1, -1, 1], [-1, 1, -1, 1], [-1, -1, 1, 1], [1, 1, -1, 1], [1, -1, 1, 1], [-1, 1, 1, 1], [1,1,1,1]])
affine64_mni = \
np.array([[ -2.85714293, -0. , 0. , 90. ],
[ -0. , 3.42857146, -0. , -126. ],
[ 0. , 0. , 2.85714293, -72. ],
[ 0. , 0. , 0. , 1. ]])
scriptpath = os.path.dirname(os.path.realpath(__file__))
device = torch.device("cpu")
net = HeadModel()
net.to(device)
net.load_state_dict(torch.load(scriptpath + "/torchparams/params_head_00075_00000.pt", map_location=device))
net.eval()
netAff = ModelAff()
netAff.load_state_dict(torch.load(scriptpath + "/torchparams/paramsaffineta_00079_00000.pt", map_location=device), strict=False)
netAff.to(device)
netAff.eval()
class HippoModel(nn.Module):
def __init__(self):
super(HippoModel, self).__init__()
self.conv0a_0 = l = nn.Conv3d(1, 16, (1,1,3), padding=0)
self.conv0a_1 = l = nn.Conv3d(16, 16, (1,3,1), padding=0)
self.conv0a = nn.Conv3d(16, 16, (3,1,1), padding=0)
self.convf1 = nn.Conv3d(16, 48, (3,3,3), padding=0)
self.maxpool1 = nn.MaxPool3d(2)
self.bn1 = nn.BatchNorm3d(48, momentum=1)
self.bn1.training = False
self.convout0 = nn.Conv3d(48, 48, (3,3,3), padding=1)
self.convout1 = nn.Conv3d(48, 48, (3,3,3), padding=1)
self.maxpool2 = nn.MaxPool3d(2)
self.bn2 = nn.BatchNorm3d(48, momentum=1)
self.bn2.training = False
self.convout2p = nn.Conv3d(48, 48, (3,3,3), padding=1)
self.convout2 = nn.Conv3d(48, 48, (3,3,3), padding=1)
self.convlx3 = nn.Conv3d(48, 48, (3,3,3), padding=1)
self.convlx5 = nn.Conv3d(48, 48, (3,3,3), padding=1)
self.convlx7 = nn.Conv3d(48, 16, (3,3,3), padding=1)
self.convlx8 = nn.Conv3d(16, 1, 1, padding=0)
self.blur = nn.Conv3d(1, 1, 7, padding=3)
self.conv_extract = nn.Conv3d(48, 47, 3, padding=1)
self.convmix = nn.Conv3d(48, 16, 3, padding=1)
self.convout1x = nn.Conv3d(16, 1, 1, padding=0)
def forward(self, x):
x = F.relu(self.conv0a_0(x))
x = F.relu(self.conv0a_1(x))
x = F.relu(self.conv0a(x))
self.out_conv_f1 = x = F.relu(self.convf1(x))
self.out_maxpool1 = x = self.maxpool1(x)
x = self.bn1(x)
x = F.relu(self.convout0(x))
x = self.convout1(x)
x = x + self.out_maxpool1
x = F.relu(x)
self.out_maxpool2 = x = self.maxpool2(x)
x = self.bn2(x)
x = F.relu(self.convout2p(x))
x = self.convout2(x)
x = x + self.out_maxpool2
x = F.relu(x)
self.lx2 = F.interpolate(x, scale_factor=2, mode="nearest")
x = F.relu(self.convlx3(x))
x = F.interpolate(x, scale_factor=2, mode="nearest")
x = F.relu(self.convlx5(x))
x = F.interpolate(x, scale_factor=2, mode="nearest")
x = F.relu(self.convlx7(x))
self.out_output1 = x = torch.sigmoid(self.convlx8(x))
x = torch.sigmoid(self.blur(x))
x = x * self.out_conv_f1
x = F.leaky_relu(self.conv_extract(x))
x = torch.cat([self.out_output1, x], dim=1)
x = F.relu(self.convmix(x))
self.out_output2 = x = torch.sigmoid(self.convout1x(x))
#x = torch.cat([self.out_output2, self.out_output1], dim=1)
return x
hipponet = HippoModel()
hipponet.load_state_dict(torch.load(scriptpath + "/torchparams/hippodeep.pt"))
OUTPUT_RES64 = False
OUTPUT_NATIVE = True
OUTPUT_DEBUG = False
allsubjects_scalar_report = []
mul_homo = lambda g, Mt : g @ Mt[:3,:3].astype(np.float32) + Mt[3,:3].astype(np.float32)
def indices_unitary(dimensions, dtype):
dimensions = tuple(dimensions)
N = len(dimensions)
shape = (1,)*N
res = np.empty((N,)+dimensions, dtype=dtype)
for i, dim in enumerate(dimensions):
res[i] = np.linspace(-1, 1, dim, dtype=dtype).reshape( shape[:i] + (dim,) + shape[i+1:] )
return res
def main():
for fname in sys.argv[1:]:
if "_mask" in fname:
print("Skipping %s because the filename contains _mask in it" % fname)
continue
Ti = time.time()
try:
print("Loading image " + fname)
outfilename = fname.replace(".mnc", ".nii").replace(".mgz", ".nii").replace(".nii.gz", ".nii").replace(".nii", "_tiv.nii.gz")
img = nibabel.load(fname)
if type(img) is nibabel.nifti1.Nifti1Image:
img._affine = img.get_qform() # for ANTs compatibility
if type(img) is nibabel.Nifti1Image:
if img.header["qform_code"] == 0:
if img.header["sform_code"] == 0:
print(" *** Error: the header of this nifti file has no qform_code defined.")
print(" Fix the header manually or reconvert from the original DICOM.")
if not OUTPUT_DEBUG:
continue
if not np.allclose(img.get_sform(), img.get_qform()):
img._affine = img.get_qform() # simplify later ANTs compatibility
print("This image has an sform defined, ignoring it - work in scanner space using the qform")
except:
open(fname + ".warning.txt", "a").write("can't open the file\n")
print(" *** Error: can't open file. Skip")
continue
d = img.get_fdata(caching="unchanged", dtype=np.float32)
while len(d.shape) > 3:
print("Warning: this looks like a timeserie. Averaging it")
open(fname + ".warning.txt", "a").write("dim not 3. Averaging last dimension\n")
d = d.mean(-1)
d = (d - d.mean()) / d.std()
o1 = nibabel.orientations.io_orientation(img.affine)
o2 = np.array([[ 0., -1.], [ 1., 1.], [ 2., 1.]]) # We work in LAS space (same as the mni_icbm152 template)
trn = nibabel.orientations.ornt_transform(o1, o2) # o1 to o2 (apply to o2 to obtain o1)
trn_back = nibabel.orientations.ornt_transform(o2, o1)
revaff1 = nibabel.orientations.inv_ornt_aff(trn, (1,1,1)) # mult on o1 to obtain o2
revaff1i = nibabel.orientations.inv_ornt_aff(trn_back, (1,1,1)) # mult on o2 to obtain o1
aff_orig64 = np.linalg.lstsq(bbox_world(np.identity(4), (64,64,64)), bbox_world(img.affine, img.shape[:3]), rcond=None)[0].T
voxscale_native64 = np.abs(np.linalg.det(aff_orig64))
revaff64i = nibabel.orientations.inv_ornt_aff(trn_back, (64,64,64))
aff_reor64 = np.linalg.lstsq(bbox_world(revaff64i, (64,64,64)), bbox_world(img.affine, img.shape[:3]), rcond=None)[0].T
wgridt = (netAff.grid @ torch.tensor(revaff1i, device=device, dtype=torch.float32))[None,...,[2,1,0]]
d_orr = F.grid_sample(torch.as_tensor(d, dtype=torch.float32, device=device)[None,None], wgridt, align_corners=True)
if OUTPUT_DEBUG:
nibabel.Nifti1Image(np.asarray(d_orr[0,0].cpu()), aff_reor64).to_filename(outfilename.replace("_tiv", "_orig_b64"))
## Head priors
T = time.time()
with torch.no_grad():
out1t = net(d_orr)
out1 = np.asarray(out1t.cpu())
#print("Head Inference in ", time.time() - T)
## Output head priors
scalar_output = []
scalar_output_report = []
# brain mask
output = out1[0,0].astype("float32")
out_cc, lab = scipy.ndimage.label(output > .01)
#output *= (out_cc == np.bincount(out_cc.flat)[1:].argmax()+1)
brainmask_cc = torch.tensor(output)
vol = (output[output > .5]).sum() * voxscale_native64
if OUTPUT_DEBUG:
print(" Estimated intra-cranial volume (mm^3): %d" % vol)
if 0:
open(outfilename.replace("_tiv.nii.gz", "_eTIV.txt"), "w").write("%d\n" % vol)
scalar_output.append(vol)
scalar_output_report.append(vol)
if OUTPUT_RES64:
out = (output.clip(0, 1) * 255).astype("uint8")
nibabel.Nifti1Image(out, aff_reor64, img.header).to_filename(outfilename.replace("_tiv", "_tissues%d_b64" % 0))
if OUTPUT_NATIVE:
# wgridt for native space
gsx, gsy, gsz = img.shape[:3]
# this is a big array, so use float16
sgrid = np.rollaxis(indices_unitary((gsx,gsy,gsz), dtype=np.float16),0,4)
wgridt = torch.as_tensor(mul_homo(sgrid, inv(revaff1i))[None,...,[2,1,0]], device=device, dtype=torch.float32)
del sgrid
dnat = np.asarray(F.grid_sample(torch.as_tensor(output, dtype=torch.float32, device=device)[None,None], wgridt, align_corners=True).cpu())[0,0]
#nibabel.Nifti1Image(dnat, img.affine).to_filename(outfilename.replace("_tiv", "_tissues%d" % 0))
nibabel.Nifti1Image((dnat > .5).astype("uint8"), img.affine).to_filename(outfilename.replace("_tiv", "_brain_mask"))
vol = (dnat > .5).sum() * np.abs(np.linalg.det(img.affine))
print(" Estimated intra-cranial volume (mm^3) (native space): %d" % vol)
scalar_output.append(vol)
scalar_output_report[-1] = vol # authoritative, so overwrite previous
del dnat
if 1:
# cerebrum mask
output = out1[0,2].astype("float32")
out_cc, lab = scipy.ndimage.label(output > .01)
output *= (out_cc == np.bincount(out_cc.flat)[1:].argmax()+1)
vol = (output[output > .5]).sum() * voxscale_native64
if OUTPUT_DEBUG:
print(" Estimated cerebrum volume (mm^3): %d" % vol)
if 0:
open(outfilename.replace("_tiv.nii.gz", "_eTIV_nocerebellum.txt"), "w").write("%d\n" % vol)
scalar_output.append(vol)
if OUTPUT_RES64:
out = (output.clip(0, 1) * 255).astype("uint8")
nibabel.Nifti1Image(out, aff_reor64, img.header).to_filename(outfilename.replace("_tiv", "_tissues%d_b64" % 2))
if OUTPUT_NATIVE:
dnat = np.asarray(F.grid_sample(torch.as_tensor(output, dtype=torch.float32, device=device)[None,None], wgridt, align_corners=True).cpu()[0,0])
#nibabel.Nifti1Image(dnat, img.affine).to_filename(outfilename.replace("_tiv", "_tissues%d" % 2))
nibabel.Nifti1Image((dnat > .5).astype("uint8"), img.affine).to_filename(outfilename.replace("_tiv", "_cerebrum_mask"))
vol = (dnat > .5).sum() * np.abs(np.linalg.det(img.affine))
print(" Estimated cerebrum volume (mm^3) (native space): %d" % vol)
scalar_output.append(vol)
del dnat
# cortex
output = out1[0,1].astype("float32")
output[output < .01] = 0
if OUTPUT_RES64:
out = (output.clip(0, 1) * 255).astype("uint8")
nibabel.Nifti1Image(out, aff_reor64, img.header).to_filename(outfilename.replace("_tiv", "_tissues%d_b64" % 1))
if OUTPUT_NATIVE and OUTPUT_DEBUG:
dnat = np.asarray(F.grid_sample(torch.as_tensor(output, dtype=torch.float32, device=device)[None,None], wgridt, align_corners=True).cpu()[0,0])
nibabel.Nifti1Image(dnat, img.affine).to_filename(outfilename.replace("_tiv", "_tissues%d" % 1))
del dnat
## MNI affine
T = time.time()
with torch.no_grad():
wc1, tA = netAff(out1t[:,[1,3]] * brainmask_cc)
wnat = np.linalg.lstsq(bbox_world(img.affine, img.shape[:3]), bbox_one @ revaff1, rcond=None)[0]
wmni = np.linalg.lstsq(bbox_world(affine64_mni, (64,64,64)), bbox_one, rcond=None)[0]
M = (wnat @ inv(np.asarray(tA[0].cpu())) @ inv(wmni)).T
# [native world coord] @ M.T -> [mni world coord] , in LAS space
if OUTPUT_DEBUG:
# Output MNI, mostly for debug, save in box64, uint8
out2 = np.asarray(wc1.to("cpu"))
out2 = np.clip((out2 * 255), 0, 255).astype("uint8")
nibabel.Nifti1Image(out2[0,0], affine64_mni).to_filename(outfilename.replace("_tiv", "_mniwrapc1"))
del out2
if 0:
out2r = np.asarray(netAff.resample_other(d_orr).cpu())
out2r = (out2r - out2r.min()) * 255 / out2r.ptp()
nibabel.Nifti1Image(out2r[0,0].astype("uint8"), affine64_mni).to_filename(outfilename.replace("_tiv", "_mniwrap"))
del out2r
# output an ANTs-compatible matrix (AntsApplyTransforms -t)
f3 = np.array([[1, 1, -1, -1],[1, 1, -1, -1], [-1, -1, 1, 1], [1, 1, 1, 1]]) # ANTs LPS
MI = inv(M) * f3
txt = """#Insight Transform File V1.0\nTransform: AffineTransform_float_3_3\nFixedParameters: 0 0 0\nParameters: """
txt += " ".join(["%4.6f %4.6f %4.6f" % tuple(x) for x in MI[:3,:3].tolist()]) + " %4.6f %4.6f %4.6f\n" % (MI[0,3], MI[1,3], MI[2,3])
if 0:
open(outfilename.replace("_tiv.nii.gz", "_mni0Affine.txt"), "w").write(txt)
u, s, vt = np.linalg.svd(MI[:3,:3])
MI3rigid = u @ vt
txt = """#Insight Transform File V1.0\nTransform: AffineTransform_float_3_3\nFixedParameters: 0 0 0\nParameters: """
txt += " ".join(["%4.6f %4.6f %4.6f" % tuple(x) for x in MI3rigid.tolist()]) + " %4.6f %4.6f %4.6f\n" % (MI[0,3], MI[1,3], MI[2,3])
if 0:
open(outfilename.replace("_tiv.nii.gz", "_mni0Rigid.txt"), "w").write(txt)
## Hippodeep
T = time.time()
imgcroproi_affine = np.array([[ -1., -0., 0., 54.], [ -0., 1., -0., -59.], [0., 0., 1., -45.], [0., 0., 0., 1.]])
imgcroproi_shape = (107, 72, 68)
# coord in mm bbox
gsx, gsy, gsz = 107, 72, 68
sgrid = np.rollaxis(indices_unitary((gsx,gsy,gsz), dtype=np.float32),0,4)
bboxnat = bbox_world(imgcroproi_affine, imgcroproi_shape) @ inv(M.T) @ wnat
matzoom = np.linalg.lstsq(bbox_one, bboxnat, rcond=None)[0] # in -1..1 space
# wgridt for hippo box
wgridt = torch.tensor(mul_homo( sgrid, (matzoom @ revaff1i) )[None,...,[2,1,0]], device=device, dtype=torch.float32)
del sgrid
dout = F.grid_sample(torch.as_tensor(d, dtype=torch.float32, device=device)[None,None], wgridt, align_corners=True)
# note: d was normalized from full-image
d_in = np.asarray(dout[0,0].cpu()) # back to numpy since torch does not support negative step/strides
if OUTPUT_RES64:
d_in_u8 = (((d_in - d_in.min()) / d_in.ptp()) * 255).astype("uint8")
nibabel.Nifti1Image(d_in_u8, imgcroproi_affine).to_filename(outfilename.replace("_tiv", "_affcrop"))
d_in -= d_in.mean()
d_in /= d_in.std()
# split Left and Right (flipping Right)
with torch.no_grad():
hippoR = hipponet(torch.as_tensor(d_in[None, None, 6: 54:+1,: ,2:-2 ].copy()))
hippoL = hipponet(torch.as_tensor(d_in[None, None,-7:-55:-1,: ,2:-2 ].copy()))
hippoRL = np.vstack([np.asarray(hippoR.cpu()), np.asarray(hippoL.cpu())])
#print("Hippo Inferrence in " + str(time.time() - T))
# smoothly rescale (.5 ~ .75) to (.5 ~ 1.)
hippoRL = np.clip(((hippoRL - .5) * 2 + .5), 0, 1) * (hippoRL > .5)
# lots numpy/torch copy below, because torch raises errors on negative strides
output = np.zeros((2, 107, 72, 68), np.float32)
output[0, -7:-55:-1,: ,2:-2][2:-2,2:-2,2:-2] = np.clip(hippoRL[1] * 255, 0, 255)#* maskL
output[1, 6: 54:+1,: ,2:-2][2:-2,2:-2,2:-2] = np.clip(hippoRL[0] * 255, 0, 255) # * maskR
if OUTPUT_DEBUG:
#outputfn = outfilename.replace(".nii.gz", "_outseg_L.nii.gz")
#nibabel.Nifti1Image(output[0], imgcroproi_affine).to_filename(outputfn)
#outputfn = outfilename.replace(".nii.gz", "_outseg_R.nii.gz")
#nibabel.Nifti1Image(output[1], imgcroproi_affine).to_filename(outputfn)
outputfn = outfilename.replace("_tiv", "_affcrop_outseg_mask")
nibabel.Nifti1Image(output.sum(0), imgcroproi_affine).to_filename(outputfn)
boxvols = hippoRL[[1,0]].reshape(2, -1).sum(1) * np.abs(np.linalg.det(imgcroproi_affine @ inv(M)))
scalar_output.append(boxvols)
if 1:
def bbox_xyz(shape, affine):
" returns the worldspace of the edge of the image "
s = shape[0]-1, shape[1]-1, shape[2]-1
bbox = [[0,0,0], [s[0],0,0], [0,s[1],0], [0,0,s[2]], [s[0],s[1],0], [s[0],0,s[2]], [0,s[1],s[2]], [s[0],s[1],s[2]]]
return mul_homo(bbox, affine.T)
def indices_xyz(shape, affine, offset_vox= np.array([0,0,0])):
assert (len(shape) == 3)
ind = np.indices(shape).astype(np.float32) + offset_vox.reshape(3, 1,1,1).astype(np.float32)
return mul_homo(np.rollaxis(ind, 0, 4), affine.T)
def xyz_to_DHW3(xyz, iaffine, srcshape):
affine = np.linalg.inv(iaffine)
ijk3 = mul_homo(xyz, affine.T)
ijk3[...,0] /= srcshape[0] -1
ijk3[...,1] /= srcshape[1] -1
ijk3[...,2] /= srcshape[2] -1
ijk3 = ijk3 * 2 - 1
DHW3 = np.swapaxes(ijk3, 0, 2)
return DHW3
pts = bbox_xyz(imgcroproi_shape, imgcroproi_affine)
pts = mul_homo(pts, np.linalg.inv(M).T)
pts_ijk = mul_homo(pts, np.linalg.inv(img.affine).T)
for i in range(3):
np.clip(pts_ijk[:,i], 0, img.shape[i], out = pts_ijk[:,i])
pmin = np.floor(np.min(pts_ijk, 0)).astype(int)
pwidth = np.ceil(np.max(pts_ijk, 0)).astype(int) - pmin
widx = indices_xyz(pwidth, img.affine, offset_vox=pmin)
widx = mul_homo(widx, M.T)
DHW3 = xyz_to_DHW3(widx, imgcroproi_affine, imgcroproi_shape)
wdata = np.zeros(img.shape[:3], np.uint8)
d = torch.tensor(output[0].T, dtype=torch.float32)
outDHW = F.grid_sample(d[None,None], torch.tensor(DHW3[None]), align_corners=True)
dnat = np.asarray(outDHW[0,0].permute(2,1,0))
dnat[dnat < 32] = 0 # remove noise
volsAA_L = dnat.sum() / 255. * np.abs(np.linalg.det(img.affine))
wdata[pmin[0]:pmin[0]+pwidth[0], pmin[1]:pmin[1]+pwidth[1], pmin[2]:pmin[2]+pwidth[2]] = dnat.astype(np.uint8)
nibabel.Nifti1Image(wdata.astype("uint8"), img.affine).to_filename(outfilename.replace("_tiv", "_mask_L"))
d = torch.tensor(output[1].T, dtype=torch.float32)
outDHW = F.grid_sample(d[None,None], torch.tensor(DHW3[None]), align_corners=True)
dnat = np.asarray(outDHW[0,0].permute(2,1,0))
dnat[dnat < 32] = 0 # remove noise
volsAA_R = dnat.sum() / 255. * np.abs(np.linalg.det(img.affine))
wdata[pmin[0]:pmin[0]+pwidth[0], pmin[1]:pmin[1]+pwidth[1], pmin[2]:pmin[2]+pwidth[2]] = dnat.astype(np.uint8)
nibabel.Nifti1Image(wdata.astype("uint8"), img.affine).to_filename(outfilename.replace("_tiv", "_mask_R"))
print(" Hippocampal volumes (L,R)", volsAA_L, volsAA_R)
scalar_output.append([volsAA_L, volsAA_R])
scalar_output_report.append([volsAA_L, volsAA_R])
if OUTPUT_DEBUG:
txt = "eTIV_mni,eTIV,cerebrum_mni,cerebrum,mni_hippoL,mni_hippoR,hippoL,hippoR\n"
txt += "%4f,%4f,%4f,%4f,%4.4f,%4.4f,%4.4f,%4.4f\n" % (tuple(scalar_output[:4]) + tuple(scalar_output[4])+ tuple(scalar_output[5]))
open(outfilename.replace("_tiv.nii.gz", "_scalars_hippo.csv"), "w").write(txt)
if 1:
txt = "eTIV,hippoL,hippoR\n"
txt += "%4f,%4f,%4f\n" % (scalar_output_report[0], scalar_output_report[1][0], scalar_output_report[1][1])
open(outfilename.replace("_tiv.nii.gz", "_hippoLR_volumes.csv"), "w").write(txt)
if OUTPUT_RES64:
print("fslview %s %s -t .5 &" % (outfilename.replace("_tiv", "_affcrop"), outfilename.replace("_tiv", "_affcrop_outseg_mask")))
print(" Elapsed time for subject %4.2fs " % (time.time() - Ti))
print(" To display using fsleyes or fslview, try:")
print(" fsleyes %s %s -a 75 -cm Red-Yellow %s -a 75 -cm Blue-Lightblue &" % (fname, outfilename.replace("_tiv", "_mask_L"), outfilename.replace("_tiv", "_mask_R")))
print(" fslview %s %s -t .5 %s -t .5 &" % (fname, outfilename.replace("_tiv", "_mask_L"), outfilename.replace("_tiv", "_mask_R")))
allsubjects_scalar_report.append( (fname, scalar_output_report[0], scalar_output_report[1][0], scalar_output_report[1][1]) )
try:
print("Peak memory used (Gb) " + str(resource.getrusage(resource.RUSAGE_SELF)[2] / (1024.*1024)))
except:
pass
print("Done")
if len(sys.argv[1:]) > 1:
outfilename = (os.path.dirname(fname) or ".") + "/all_subjects_hippo_report.csv"
txt_entries = ["%s,%4f,%4f,%4f\n" % s for s in allsubjects_scalar_report]
open(outfilename, "w").writelines( [ "filename,eTIV,hippoL,hippoR\n" ] + txt_entries)
print("Volumes of every subjects saved as " + outfilename)
if __name__ == "__main__":
main()
