invertcontrast.py

import ismrmrd
import os
import itertools
import logging
import numpy as np
import numpy.fft as fft
import base64

# Folder for debug output files
debugFolder = "/tmp/share/debug"

def process(connection, config, metadata):
    logging.info("Config: \n%s", config)

    # Metadata should be MRD formatted header, but may be a string
    # if it failed conversion earlier
    try:
        logging.info("Metadata: \n%s", metadata.toxml('utf-8'))

        logging.info("Incoming dataset contains %d encodings", len(metadata.encoding))
        logging.info("First encoding is of type '%s', with a field of view of (%s x %s x %s)mm^3 and a matrix size of (%s x %s x %s)", 
            metadata.encoding[0].trajectory, 
            metadata.encoding[0].encodedSpace.matrixSize.x, 
            metadata.encoding[0].encodedSpace.matrixSize.y, 
            metadata.encoding[0].encodedSpace.matrixSize.z, 
            metadata.encoding[0].encodedSpace.fieldOfView_mm.x, 
            metadata.encoding[0].encodedSpace.fieldOfView_mm.y, 
            metadata.encoding[0].encodedSpace.fieldOfView_mm.z)

    except:
        logging.info("Improperly formatted metadata: \n%s", metadata)

    # Continuously parse incoming data parsed from MRD messages
    acqGroup = []
    imgGroup = []
    waveformGroup = []
    try:
        for item in connection:
            # ----------------------------------------------------------
            # Raw k-space data messages
            # ----------------------------------------------------------
            if isinstance(item, ismrmrd.Acquisition):
                # Accumulate all imaging readouts in a group
                if (not item.is_flag_set(ismrmrd.ACQ_IS_NOISE_MEASUREMENT) and
                    not item.is_flag_set(ismrmrd.ACQ_IS_PHASECORR_DATA)):
                    acqGroup.append(item)

                # When this criteria is met, run process_raw() on the accumulated
                # data, which returns images that are sent back to the client.
                if item.is_flag_set(ismrmrd.ACQ_LAST_IN_SLICE):
                    logging.info("Processing a group of k-space data")
                    image = process_raw(acqGroup, config, metadata)
                    logging.debug("Sending image to client:\n%s", image)
                    connection.send_image(image)
                    acqGroup = []

            # ----------------------------------------------------------
            # Image data messages
            # ----------------------------------------------------------
            elif isinstance(item, ismrmrd.Image):
                # Only process magnitude images -- send phase images back without modification
                if item.image_type is ismrmrd.IMTYPE_MAGNITUDE:
                    imgGroup.append(item)
                else:
                    connection.send_image(item)
                    continue

                # When this criteria is met, run process_group() on the accumulated
                # data, which returns images that are sent back to the client.
                # TODO: logic for grouping images
                if False:
                    logging.info("Processing a group of images")
                    image = process_image(imgGroup, config, metadata)
                    logging.debug("Sending image to client:\n%s", image)
                    connection.send_image(image)
                    imgGroup = []

            # ----------------------------------------------------------
            # Waveform data messages
            # ----------------------------------------------------------
            elif isinstance(item, ismrmrd.Waveform):
                waveformGroup.append(item)

            elif item is None:
                break

            else:
                logging.error("Unsupported data type %s", type(item).__name__)

        # Extract raw ECG waveform data. Basic sorting to make sure that data 
        # is time-ordered, but no additional checking for missing data.
        # ecgData has shape (5 x timepoints)
        if len(waveformGroup) > 0:
            waveformGroup.sort(key = lambda item: item.time_stamp)
            ecgData = [item.data for item in waveformGroup if item.waveform_id == 0]
            ecgData = np.concatenate(ecgData,1)

        # Process any remaining groups of raw or image data.  This can 
        # happen if the trigger condition for these groups are not met.
        # This is also a fallback for handling image data, as the last
        # image in a series is typically not separately flagged.
        if len(acqGroup) > 0:
            logging.info("Processing a group of k-space data (untriggered)")
            image = process_raw(acqGroup, config, metadata)
            logging.debug("Sending image to client:\n%s", image)
            connection.send_image(image)
            acqGroup = []

        if len(imgGroup) > 0:
            logging.info("Processing a group of images (untriggered)")
            image = process_image(imgGroup, config, metadata)
            logging.debug("Sending image to client:\n%s", image)
            connection.send_image(image)
            imgGroup = []

    finally:
        connection.send_close()


def process_raw(group, config, metadata):
    # Create folder, if necessary
    if not os.path.exists(debugFolder):
        os.makedirs(debugFolder)
        logging.debug("Created folder " + debugFolder + " for debug output files")

    # Sort by line number (incoming data may be interleaved)
    lin = [acquisition.idx.kspace_encode_step_1 for acquisition in group]
    logging.debug("Incoming lin ordering: " + str(lin))

    group.sort(key = lambda acq: acq.idx.kspace_encode_step_1)
    sortedLin = [acquisition.idx.kspace_encode_step_1 for acquisition in group]
    logging.debug("Sorted lin ordering: " + str(sortedLin))

    # Format data into single [cha RO PE] array
    data = [acquisition.data for acquisition in group]
    data = np.stack(data, axis=-1)

    logging.debug("Raw data is size %s" % (data.shape,))
    np.save(debugFolder + "/" + "raw.npy", data)

    # Fourier Transform
    data = fft.fftshift(data, axes=(1, 2))
    data = fft.ifft2(data)
    data = fft.ifftshift(data, axes=(1, 2))

    # Sum of squares coil combination
    data = np.abs(data)
    data = np.square(data)
    data = np.sum(data, axis=0)
    data = np.sqrt(data)

    logging.debug("Image data is size %s" % (data.shape,))
    np.save(debugFolder + "/" + "img.npy", data)

    # Normalize and convert to int16
    data *= 32767/data.max()
    data = np.around(data)
    data = data.astype(np.int16)

    # Remove phase oversampling
    nRO = np.size(data,0)
    data = data[int(nRO/4):int(nRO*3/4),:]
    logging.debug("Image without oversampling is size %s" % (data.shape,))
    np.save(debugFolder + "/" + "imgCrop.npy", data)

    # Format as ISMRMRD image data
    image = ismrmrd.Image.from_array(data, acquisition=group[0])
    image.image_index = 1

    # Set ISMRMRD Meta Attributes
    meta = ismrmrd.Meta({'DataRole':               'Image',
                         'ImageProcessingHistory': ['FIRE', 'PYTHON'],
                         'WindowCenter':           '16384',
                         'WindowWidth':            '32768'})
    xml = meta.serialize()
    logging.debug("Image MetaAttributes: %s", xml)
    logging.debug("Image data has %d elements", image.data.size)

    image.attribute_string = xml

    # Call process_image() to invert image contrast
    image = process_image([image], config, metadata)

    return image


def process_image(images, config, metadata):
    # Create folder, if necessary
    if not os.path.exists(debugFolder):
        os.makedirs(debugFolder)
        logging.debug("Created folder " + debugFolder + " for debug output files")

    logging.debug("Incoming image data of type %s", ismrmrd.get_dtype_from_data_type(images[0].data_type))

    # Display MetaAttributes for first image
    meta = ismrmrd.Meta.deserialize(images[0].attribute_string)
    logging.debug("MetaAttributes: %s", ismrmrd.Meta.serialize(meta))

    # Optional serialization of ICE MiniHeader
    if 'IceMiniHead' in meta:
        logging.debug("IceMiniHead: %s", base64.b64decode(meta['IceMiniHead']).decode('utf-8'))

    # Extract image data into a 5D array of size [img cha z y x]
    data = np.stack([img.data for img in images])
    head = [img.getHead() for img in images]

    logging.debug("Original image data is size %s" % (data.shape,))
    np.save(debugFolder + "/" + "imgOrig.npy", data)

    # Normalize and convert to int16
    data = data.astype(np.float64)
    data *= 32767/data.max()
    data = np.around(data)
    data = data.astype(np.int16)

    # Invert image contrast
    data = 32767-data
    data = np.abs(data)
    data = data.astype(np.int16)
    np.save(debugFolder + "/" + "imgInverted.npy", data)

    # Re-slice back into 2D images
    imagesOut = [None] * data.shape[0]
    for iImg in range(data.shape[0]):
        # Create new MRD instance for the inverted image
        imagesOut[iImg] = ismrmrd.Image.from_array(data[iImg,...].transpose())
        data_type = imagesOut[iImg].data_type

        # Copy the fixed header information
        oldHeader = head[iImg]
        oldHeader.data_type = data_type

        imagesOut[iImg].setHead(oldHeader)

        # Set ISMRMRD Meta Attributes
        meta = ismrmrd.Meta({'DataRole':               'Image',
                            'ImageProcessingHistory': ['FIRE', 'PYTHON'],
                            'WindowCenter':           '16384',
                            'WindowWidth':            '32768'})
        xml = meta.serialize()
        logging.debug("Image MetaAttributes: %s", xml)
        logging.debug("Image data has %d elements", imagesOut[iImg].data.size)

        imagesOut[iImg].attribute_string = xml

    return imagesOut