Source code for datasets.brain.serialize

import glob
import os

import nibabel as nib
import nilearn as nl
import tensorflow as tf
import numpy as np
import pandas as pd

from ingredient_wrapper import Ingredient
from sacred import Experiment
from handlers import NiftiHandler, Int64Handler, Float32Handler
from datasets.serialize import serializer_ingred, serialize_metacsv
from .get_raw import raw_brain_ingred

brain_serializer_ingred = Ingredient("brain_serializer",
                                     ingredients=[serializer_ingred,
                                                  raw_brain_ingred])

@brain_serializer_ingred.config
def serializer_updates():
    record_dir = "datasets/brain/records"
    record_pattern = "brain_{split}_{idx}.tfrecord"
    samples_per_record = 32
    split_to_size = {
        "train": 0.6,
        "eval": 0.2,
        "test": 0.2
    }
    compression = "GZIP"


brain3D_serializer_ingred = Ingredient("brain3D_serializer",
    ingredients = [brain_serializer_ingred])

@brain3D_serializer_ingred.config
def config(raw_dir):

    img_shape = [128, 128, 128]
    img_dtype = "float32"

    keys_to_descriptions = {"image": "Whole 3D Brain Image",
                            "image/shape": "Shape of 3D Brain Image",
                            "age": "Subject Age"}

    # TODO: Turn classes into serialized description and parse later
    keys_to_handlers = {
        "image": NiftiHandler(raw_dir, img_shape, img_dtype),
        "image/shape": Int64Handler(shape=[3], delegate_to="image"),
        "age":   Float32Handler()
    }


brain2D_serializer_ingred = Ingredient("brain2D_serializer",
    ingredients = [brain_serializer_ingred])

@brain2D_serializer_ingred.config
def config(raw_dir):

    img_shape = [128, 128]
    img_dtype = "float32"
    img_slice = (0,64)

    keys_to_descriptions = {"image": "2D Brain Slice",
                            "image/shape": "Shape of 2D Brain Slice",
                            "age": "Subject Age"}

    # TODO: Turn classes into serialized description and parse later
    keys_to_handlers = {
        "image": NiftiHandler(raw_dir, img_shape, img_dtype, img_slice),
        "image/shape": Int64Handler(shape=[2], delegate_to="image"),
        "age":   Float32Handler()
    }


@brain_serializer_ingred.command
def rename(raw_dir):
    for file in glob.glob(os.path.join(raw_dir, "*")):
        if "nii.gz" in file:
            idx = file.split("/")[-1].split("_")[0]
            newfile = os.path.join(raw_dir, idx + ".nii.gz")
            os.rename(file, newfile)


@brain2D_serializer_ingred.command
@brain3D_serializer_ingred.command
def serialize(csv_file, keys_to_handlers):

    serialize_metacsv(csv_file=csv_file)


if __name__ == '__main__':
    """
    CommandLine Usage:
    python -m datasets.brain.serialize with brain_serializer.record_dir=datasets/brain/records brain_serializer.raw_dir=datasets/brain/test_data brain_serializer.csv_file=datasets/brain/test_data/meta_data.csv
    """
    ex = Experiment("Serialize_Brain", ingredients=[brain_serializer_ingred])

    @ex.main
    def main():
      serialize()

    ex.run_commandline()