eval_finetuned.py

import matplotlib.pyplot as plt
import os
import numpy as np
import glob
from matplotlib import gridspec
import numpy as np
from PIL import Image
import tensorflow as tf
from tensorflow.python.platform import gfile
import glob
import argparse
import cv2
import nibabel as nib
from tqdm import tqdm

# adapted from https://colab.research.google.com/github/tensorflow/models/blob/master/research/deeplab/deeplab_demo.ipynb


class DeepLabModel(object):
  """Class to load deeplab model and run inference."""

  FEATURES_TENSOR_NAME=  'import/decoder/decoder_conv1_pointwise/Conv2D:0' 
  INPUT_TENSOR_NAME = 'import/ImageTensor:0'
  OUTPUT_TENSOR_NAME = 'import/SemanticPredictions:0'
  FROZEN_GRAPH_NAME = 'frozen_inference_graph'


  def __init__(self, name):
    """Creates and loads pretrained deeplab model."""

    self.name = name
    self.config = tf.ConfigProto(allow_soft_placement = True, gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.4))
    self.sess = tf.Session(config = self.config) 
    with gfile.FastGFile(name, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        g_in = tf.import_graph_def(graph_def)

    self.graph = tf.get_default_graph()

  def run(self, image):
    """Runs inference on a single image.

    Args:
      image: A PIL.Image object, raw input image.

    Returns:
      resized_image: RGB image resized from original input image.
      seg_map: Segmentation map of `resized_image`.
    """
    width, height = image.size
    
    image = image.convert('RGB') 
    batch_seg_map = self.sess.run(
        self.OUTPUT_TENSOR_NAME,
        feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(image)]})
    seg_map = batch_seg_map[0]
    return image, seg_map



def create_pascal_label_colormap():
    """Creates a label colormap used in PASCAL VOC segmentation benchmark.

    Returns:
     A Colormap for visualizing segmentation results.
    """
    colormap = np.zeros((256, 3), dtype=int)
    ind = np.arange(256, dtype=int)

    for shift in reversed(range(8)):
     for channel in range(3):
       colormap[:, channel] |= ((ind >> channel) & 1) << shift
     ind >>= 3
    return colormap


def label_to_color_image(label):
    """Adds color defined by the dataset colormap to the label.

    Args:
     label: A 2D array with integer type, storing the segmentation label.

    Returns:
     result: A 2D array with floating type. The element of the array
       is the color indexed by the corresponding element in the input label
       to the PASCAL color map.

    Raises:
     ValueError: If label is not of rank 2 or its value is larger than color
       map maximum entry.
    """
    if label.ndim != 2:
     raise ValueError('Expect 2-D input label')

    colormap = create_pascal_label_colormap()

    if np.max(label) >= len(colormap):
     raise ValueError('label value too large.')

    return colormap[label]


def viz_segmentation(image, seg_map, name):
    """Visualizes input image, segmentation map and overlay view."""

    plt.figure(figsize=(15, 5))
    grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1])

    plt.clf()
    plt.subplot(grid_spec[0])
    plt.imshow(image)
    plt.axis('off')
    plt.title('input image')

    plt.subplot(grid_spec[1])
    seg_image = label_to_color_image(seg_map).astype(np.uint8)
    plt.imshow(seg_image)
    plt.axis('off')
    plt.title('segmentation map')


    plt.subplot(grid_spec[2])
    plt.imshow(image)
    plt.imshow(seg_image, alpha=0.7)
    plt.axis('off')
    plt.title('segmentation overlay')
    plt.savefig('output/%s_seg_overlay_finetuned.png'%(name))
    plt.close()




   

     

def main():

    assert not ((args.nifti is not None) and (args.img_dir is not None)), 'must specify EITHER a NIfTI input file OR a directory with 2d images'
    assert (args.nifti is not None) or (args.img_dir is not None), 'must specify EITHER a NIfTI input file OR a directory with 2d images'
    if args.nifti is not None:
        NIFTI = True
        IMG = False
    else:
        IMG = True
        NIFTI = True
   
    print("writing output segmentation images to 'output'")
    os.system('mkdir -p output') 

    # get output file name
    if args.name is None:
        if args.img_dir is not None:
            out_name = args.img_dir[: args.img_dir.rfind('/')]
            out_name = out_name[out_name.rfind('/')+1:]
        else:
            out_name = args.nifti[args.nifti.rfind('/')+1:-4]
    else:
        out_name = args.name

    model = DeepLabModel('models/finetuned_model_graph.pb') 

    if IMG:     
        files = glob.glob('{img_dir}/*png'.format(img_dir=args.img_dir)) 
        assert len(files) > 0 , 'img dir must have at least one img file' 
    else:
        # load nifti
        img = nib.load(args.nifti)
        img = img.get_fdata().transpose(0, 2, 1) 
        # scale to [0, 255]
        img = 255 * img / (np.max(img) - np.min(img))
        

    segs = []
    if IMG:
        files = sorted(files)
        for f in tqdm(files, desc='segmenting images'):
            original_im = Image.open(f) 
            resized_im, seg_map  = model.run(original_im)
            resized_im = np.array(resized_im.resize((256, 256), resample=Image.BILINEAR))
            seg_map = cv2.resize(seg_map, dsize=(256, 256), interpolation=cv2.INTER_NEAREST) 
            segs.append(seg_map[None, ...])
            fname = f[f.rfind('/')+1:-4]
            if args.viz: viz_segmentation(resized_im, seg_map, fname)

    elif NIFTI:
        segs= []
        for i in tqdm(range(img.shape[0]), desc='segmenting NIfTI sections'):
            im = img[i]
            im = np.flipud(im)
            im = Image.fromarray(im)
            resized_im, seg_map  = model.run(im) 
            resized_im = np.array(resized_im.resize((256, 256), resample=Image.BILINEAR))
            seg_map = cv2.resize(seg_map, dsize=(256, 256), interpolation=cv2.INTER_NEAREST) 
            segs.append(seg_map[None, ...])
            if args.viz: viz_segmentation(resized_im, seg_map, '%s_%s' %(out_name, i))

    # concatenating 2d segmentations to 3d output
    seg = np.concatenate(segs)
    seg = nib.Nifti1Image(seg, affine=np.eye(4)) 
    nib.save(seg, 'output/{out_name}.nii.gz'.format(out_name=out_name)) 

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--nifti', type=str, help='MRI 3d input in NIfTI file format')
    parser.add_argument('--img_dir', type=str, help='directory with 2d MRI images in png format')
    parser.add_argument('--viz', type=int, default=0, help='option to visualize segmentations overlaid over MRIs') 
    parser.add_argument('--name', type=str, help='name for output segmentation file. if not specified, defaults to img_dir subdirectory name or input nifti filename')

    args = parser.parse_args()
    main()