get_mask.py

import h5py
import numpy as np
from skimage.transform import resize

import os
import logging
from os.path import join, expanduser

logger = logging.getLogger(__name__)

np.random.seed(123)


f_masks_dir = '/home/lex/mount/tdrive/Phys Research/Users/Lekan/f_masks'

def get_mask(mask, f_masks_dir=None, disp=False):
      """
      where @param `mask` = '007', for example, would represent case 7 in the f_masks directory given as a global above
            @param `f_masks_dir` is a directory given as a global above
      """
    # used to check fields in hdf5 files
    def see(n):
        print(n)

    def exp(x, axis):
        return np.expand_dims(x, axis=axis)

    with h5py.File('{}/{}.h5'.format(f_masks_dir, mask), 'r') as mf:

        if disp:
            mf.visit(see)

        masks = mf['oar_ptvs']

        logger.debug(' Loading patient {}'.format(mask))

        keys = [k for k, v in masks.items()]
        if "structs" in keys: # augmented data
            struct = masks['structs'].value

            bundle = dict(struct = struct)
        else:
            flu_grp = mf['fluence_grp']

            fluence = flu_grp['fluence'].value

            organ_names = []
            for k, _ in masks.items():
                if k != "dose":
                    organ_names.append(k)
            print(organ_names) if disp else None

            bladder     = masks['bladder'].value
            body        = masks['body'].value
            dose_mask   = masks['dose'].value
            fem_head_lt = masks['fem_head_lt'].value if 'fem_head_lt' in organ_names else masks['fem head lt'].value
            fem_head_rt = masks['fem_head_rt'].value if 'fem_head_rt' in organ_names else masks['fem head rt'].value
            ptv         = masks['ptv'].value
            rectum      = masks['rectum'].value

            struct = np.concatenate((exp(bladder, -1), exp(fem_head_lt, -1), \
                                     exp(fem_head_rt, -1), exp(body, -1),
                                     exp(ptv, -1), exp(rectum, -1)), axis=-1)
            struct = struct.transpose(-1, -2, 0, 1) # Organ x Slice x H x W

            bundle = dict(	bladder		=	bladder, 	 dose_mask	  =	dose_mask,
                            fem_head_lt = 	fem_head_lt, fem_head_rt  = fem_head_rt,
                            ptv			= 	ptv, 		 rectum		  = rectum,
                            organ_names = 	organ_names, struct       = struct,
                            fluence     =   fluence,     body         = body)

    return bundle