tools.py 1.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. import os
  2. import nibabel as nib
  3. from PIL import Image
  4. import numpy as np
  5. from sklearn.metrics import mutual_info_score
  6. import shutil
  7. from preprocess import mask_foreground
  8. def save_slice(template, filename, index, affine=np.eye(4), header=None):
  9. _slice = nib.load(str(template))
  10. _slice_data = _slice.get_fdata()
  11. newimg = nib.Nifti1Image(_slice_data[:, index, :], affine, header)
  12. newimg.header['pixdim'] = header['pixdim']
  13. if header is None:
  14. newimg.header['pixdim'][1:3] = _slice.header['pixdim'][1], _slice.header['pixdim'][3]
  15. newimg.to_filename(str(filename))
  16. return filename
  17. def create_dir(name):
  18. if not os.path.exists(name):
  19. os.mkdir(name)
  20. def remove_dir(loc):
  21. shutil.rmtree(loc)
  22. def mutual_info_mask(slice1,slice2, bins=32):
  23. slice1_masked = mask_foreground(slice1)
  24. slice2_masked = mask_foreground(slice2)
  25. common = np.logical_and(slice1_masked, slice2_masked)
  26. slice1, slice2 = np.where(common, slice1, 0), np.where(common, slice2, 0)
  27. hist = np.histogram2d(slice1.ravel(), slice2.ravel(), bins=bins)[0]
  28. return mutual_info_score(None, None, contingency=hist)
  29. def dice_coef(slice1, slice2):
  30. mask_1 = mask_foreground(slice1).astype(np.bool)
  31. mask_2 = mask_foreground(slice2).astype(np.bool)
  32. intersection = np.logical_and(mask_1, mask_2)
  33. return 2. * intersection.sum() / (mask_1.sum() + mask_2.sum())
  34. def resize_im(image_data, image_with_dimensions):
  35. dimensions = (image_with_dimensions.shape[1], image_with_dimensions.shape[0])
  36. resized = np.array(Image.fromarray(image_data).resize(dimensions))
  37. return resized
  38. def bregma_to_slice_index(bregma):
  39. return round(27.908*bregma + 116.831)