WebBasedQCImageCreator

#!/usr/bin/python

import sys
import os
import traceback
import shutil
import tempfile
import StringIO
import numpy
import PIL.Image
import dicom
import nibabel
import xnat

class GeneratorError(Exception):

    "exception for generators rejecting scans"

    def __init__(self, reason):
        self.reason = reason
        return

    def __str__(self):
        return self.reason

class Scan:

    def __init__(self):
        self._dicom_files = None
        return

    @property
    def dicom_files(self):
        if self._dicom_files is None:
            files = []
            for fname in self.files('DICOM'):
                try:
                    do = dicom.read_file(fname)
                    files.append((do.InstanceNumber, fname))
                except:
                    pass
            files.sort(lambda a, b: cmp(a[0], b[0]))
            self._dicom_files = [ f[1] for f in files ]
        return self._dicom_files

    def __repr__(self):
        return '<Scan %s>' % self.id

    def remove_copies(self):
        if self.copy_dir:
            shutil.rmtree(self.copy_dir)
            self.copy_dir = None
        return

    def save_image(self, im):
        self._reset_resource('SNAPSHOTS')
        fname = '%s_%s_qc.gif' % (self.session_label, self.id)
        self.save_single_image(im, 'SNAPSHOTS', fname, 'ORIGINAL')
        fname = '%s_%s_qc_t.gif' % (self.session_label, self.id)
        half_size = [ s/2 for s in im.size ]
        im2 = im.resize(half_size)
        self.save_single_image(im2, 'SNAPSHOTS', fname, 'THUMBNAIL')
        return

    def files(self, resource):
        if self.copy_dir is None:
            self.copy_resource(resource)
        resource_dir = '%s/%s' % (self.copy_dir, resource)
        if not os.path.exists(resource_dir):
            self.copy_resource(resource)
        for fname in os.listdir(resource_dir):
            yield '%s/%s' % (resource_dir, fname)
        return

class DirScan(Scan):

    def __init__(self, session_label, dir):
        Scan.__init__(self)
        self.session_label = session_label
        self.dir = dir
        self.id = os.path.basename(dir)
        self.resources = os.listdir(dir)
        self.copy_dir = None
        return

    def _reset_resource(self, resource):
        resource_dir = '%s/%s' % (self.dir, resource)
        if resource in self.resources:
            shutil.rmtree(resource_dir)
        os.mkdir(resource_dir)
        self.resources = os.listdir(self.dir)
        return

    def copy_resource(self, resource):
        if resource not in self.resources:
            raise ValueError('no resource %s' % resource)
        if not self.copy_dir:
            self.copy_dir = tempfile.mkdtemp()
        copy_dir = '%s/%s' % (self.copy_dir, resource)
        if os.path.exists(copy_dir):
            return
        shutil.copytree('%s/%s' % (self.dir, resource), copy_dir)
        return

    def save_single_image(self, im, resource, fname, content):
        resource_dir = '%s/%s' % (self.dir, resource)
        fo = open('%s/%s' % (resource_dir, fname), 'w')
        im.save(fo, format='GIF')
        fo.close()
        return

class XNATScan(Scan):

    def __init__(self, session_label, xnat_scan):
        Scan.__init__(self)
        self.session_label = session_label
        self.xnat_scan = xnat_scan
        self.id = self.xnat_scan.id
        self.resources = self.xnat_scan.resources
        self.copy_dir = None
        return

    def _reset_resource(self, resource):
        if resource in self.resources:
            self.resources[resource].pyxnat_resource.delete()
        self.xnat_scan.create_resource(resource)
        return

    def copy_resource(self, resource):
        if resource not in self.resources:
            raise ValueError('no resource %s' % resource)
        if not self.copy_dir:
            self.copy_dir = tempfile.mkdtemp()
        copy_dir = '%s/%s' % (self.copy_dir, resource)
        if os.path.exists(copy_dir):
            return
        os.mkdir(copy_dir)
        for file in self.resources[resource].files.itervalues():
            local_fname = '%s/%s' % (copy_dir, os.path.basename(file.path))
            file.get(local_fname)
        return

    def save_single_image(self, im, resource, fname, content):
        # resource should exist from Scan.save_image() calling 
        # XNATScan._reset_resource()
        xnat_resource = self.xnat_scan.resources[resource]
        if fname in xnat_resource.files:
            xnat_resource.files[fname].delete()
        buf = StringIO.StringIO()
        im.save(buf, format='GIF')
        xnat_resource.create_file(buf.getvalue(), 
                                  fname, 
                                  format='GIF', 
                                  content=content)
        buf.close()
        return

class TiledImage:

    """Tiled image generator

    usage:

    ti = TiledImage(n_cols, n_rows)
    ti.add_frame((0, 0), im) # im is a PIL.Image.Image instance)
    ti.add_frame((1, 0), im) # im.mode and im.shape is checked against the first
    ...
    ti.add_frame((n_x-1, n_y-1), im)
    im = ti.generate() # generates a PIL.Image.Image
    """

    # may actually be smaller if the passed frames are smaller
    target_frame_side = 280

    def __init__(self, n_cols, n_rows):
        self.tile_cols = n_cols
        self.tile_rows = n_rows
        # frames are addressed by frames[col][row]
        self.frames = []
        for c in range(self.tile_cols):
            self.frames.append(self.tile_rows * [None])
        self.ref_image = None
        return

    def __len__(self):
        "total number of frames"
        return self.tile_cols * self.tile_rows

    def add_frame(self, spec, im):
        if isinstance(spec, tuple):
            (col, row) = spec
        elif isinstance(spec, int):
            col = spec % self.tile_cols
            row = (spec-col) / self.tile_rows
        else:
            raise TypeError('spec must be a (col, row) tuple or an integer')
        if col < 0 or col >= self.tile_cols:
            raise IndexError('col out of range')
        if row < 0 or row >= self.tile_rows:
            raise IndexError('row out of range')
        if self.ref_image:
            if self.frames[col][row]:
                raise ValueError('image already set for (%d, %d)' % (col, row))
            if im.mode != self.ref_image.mode:
                msg = 'mode mismatch (%s, expecting %s)' \
                      % (im.mode, self.ref_image.mode)
                raise ValueError(msg)
            if im.size != self.ref_image.size:
                msg = 'size mismatch (%s, expecting %s)' \
                      % (str(im.size), str(self.ref_image.size))
                raise ValueError(msg)
        else:
            self.ref_image = im
        self.frames[col][row] = im
        return

    def generate(self):
        if not self.ref_image:
            raise ValueError('no images given')
        for row in range(self.tile_rows):
            for col in range(self.tile_cols):
                if not self.frames[col][row]:
                    im = PIL.Image.new(self.ref_image.mode, 
                                       self.ref_image.size, 
                                       0x00)
                    self.frames[col][row] = im
        if self.ref_image.mode == 'RGB':
            return self._generate_rgb()
        elif self.ref_image.mode == 'I':
            return self._generate_grayscale()
        raise ValueError('unsupported image mode "%s"' % self.ref_image.mode)

    def _generate_grayscale(self):
        max_vals = []
        for row in range(self.tile_rows):
            for col in range(self.tile_cols):
                max_vals.append(self.frames[col][row].getextrema()[1])
        max_val = max(max_vals)
        scale = float(255) / max_val
        native_frame_side = max(self.ref_image.size)
        frame_side = min((native_frame_side, self.target_frame_side))
        im_size = (self.tile_cols * frame_side, self.tile_rows * frame_side)
        im = PIL.Image.new('L', im_size, 0x00)
        for row in range(self.tile_rows):
            for col in range(self.tile_cols):
                square = self._square(self.frames[col][row])
                square = square.point(lambda i: i * scale, 'L')
                bbox = (col*frame_side, 
                        row*frame_side, 
                        (col+1)*frame_side, 
                        (row+1)*frame_side)
                im.paste(square.resize((frame_side, frame_side)), bbox)
        return im

    def _generate_rgb(self):
        native_frame_side = max(self.ref_image.size)
        frame_side = min((native_frame_side, self.target_frame_side))
        im_size = (self.tile_cols * frame_side, self.tile_rows * frame_side)
        im = PIL.Image.new(self.ref_image.mode, im_size, 0x00)
        for row in range(self.tile_rows):
            for col in range(self.tile_cols):
                square = self._square(self.frames[col][row])
                bbox = (col*frame_side, 
                        row*frame_side, 
                        (col+1)*frame_side, 
                        (row+1)*frame_side)
                im.paste(square.resize((frame_side, frame_side)), bbox)
        return im

    def _square(self, im):
        "pad a frame with black to return a square image"
        side = max(im.size)
        left = (side - im.size[0]) / 2
        top = (side - im.size[1]) / 2
        square_im = PIL.Image.new(im.mode, (side, side), 0x00)
        square_im.paste(im, (left, top, left+im.size[0], top+im.size[1]))
        return square_im

def command_line_error(msg):
    sys.stderr.write('%s: %s\n' % (progname, msg))
    sys.stderr.write('run %s with no arguments for usage\n' % progname)
    return

def sample(n, total):
    """sample(n, total)

    return n equally spaced integers in the range 0...total-1

    if n > total, return 0...total
    """
    if n > total:
        return range(total)
    return [ i*total/n for i in range(n) ]

def gen_clear(scan):
    return PIL.Image.fromstring('L', (256, 256), 256*256*chr(128))

def gen_dicom_survey(scan):
    if 'DICOM' not in scan.resources:
        raise GeneratorError('no DICOM resource')
    if not scan.dicom_files:
        raise GeneratorError('no DICOM files found')
    if dicom.read_file(scan.dicom_files[0]).SeriesDescription != 'Survey':
        raise GeneratorError('series description is not "Survey"')
    # we expect ImageOrientationPatient to give us images normal to 
    # primary axes
    # we bail out if any image doesn't comply, and sort the images by the 
    # index of the perpendicular axis
    images_by_axis = {0: [], 1: [], 2: []}
    for fname in scan.dicom_files:
        do = dicom.read_file(fname)
        axes = numpy.array(do.ImageOrientationPatient).reshape((2, 3))
        normal = abs(numpy.cross(axes[0], axes[1]))
        one_indices = numpy.where(normal == 1.0)[0]
        if len(one_indices) != 1:
            raise GeneratorError('oblique image found')
        images_by_axis[one_indices[0]].append(do)
    # calculate the number of images to the side of the tiled composite
    # this will be the maximum number of images in a given direction or 5, 
    # whichever is less
    n_cols = max([ len(ims) for ims in images_by_axis.itervalues() ])
    if n_cols > 5:
        n_cols = 5
    ti = TiledImage(5, 3)
    for axis in (0, 1, 2):
        images = images_by_axis[axis]
        for (tile_index, image_index) in enumerate(sample(n_cols, len(images))):
            do = images[image_index]
            data = do.pixel_array.astype('int32')
            im = PIL.Image.fromstring('I', (do.Columns, do.Rows), data)
            ti.add_frame((tile_index, axis), im)
    return ti.generate()

def gen_dicom_generic(scan):
    if 'DICOM' not in scan.resources:
        raise GeneratorError('no DICOM resource')
    if not scan.dicom_files:
        raise GeneratorError('no DICOM files found')
    ti = TiledImage(5, 5)
    for (index, slice) in enumerate(sample(len(ti), len(scan.dicom_files))):
        do = dicom.read_file(scan.dicom_files[slice])
        data = do.pixel_array.astype('int32')
        im = PIL.Image.fromstring('I', (do.Columns, do.Rows), data)
        ti.add_frame(index, im)
    return ti.generate()

def gen_dicom_rgb(scan):
    if 'DICOM' not in scan.resources:
        raise GeneratorError('no DICOM resource')
    if not scan.dicom_files:
        raise GeneratorError('no DICOM files found')
    if dicom.read_file(scan.dicom_files[0]).PhotometricInterpretation != 'RGB':
        raise GeneratorError('not RGB')
    ti = TiledImage(5, 5)
    for (index, slice) in enumerate(sample(len(ti), len(scan.dicom_files))):
        do = dicom.read_file(scan.dicom_files[slice])
        im = PIL.Image.fromstring('RGB', (do.Columns, do.Rows), do.pixel_array)
        ti.add_frame(index, im)
    return ti.generate()

def gen_nifti_generic(scan):
    if 'NIfTI' not in scan.resources:
        raise GeneratorError('no NIfTI resource')
    vol = None
    for fname in scan.files('NIfTI'):
        try:
            vol = nibabel.load(fname)
        except GeneratorError:
            raise
        except:
            pass
        if vol:
            break
    if not vol:
        raise GeneratorError('no NIfTI files found')
    if len(vol.shape) != 3:
        raise GeneratorError('non-3D NIfTI not supported')
    if vol.get_data_dtype() != 'int16':
        raise GeneratorError('non-int16 NIfTI not supported')
    ti = TiledImage(5, 5)
    for (index, slice) in enumerate(sample(len(ti), vol.shape[2])):
        slice_data = vol.get_data()[:,:,slice].astype('int32')
        im = PIL.Image.fromstring('I', vol.shape[:2], slice_data)
        ti.add_frame(index, im)
    return ti.generate()

generators = (gen_dicom_survey, 
              gen_dicom_rgb, 
              gen_dicom_generic, 
              gen_nifti_generic)

progname = os.path.basename(sys.argv.pop(0))

if not sys.argv:
    print
    print 'usage: %s <options>' % progname
    print
    print 'original options are:'
    print
    print '    -project <project that the MR-session belongs to> (ignored)'
    print '    -session <MR-session label>'
    print '    -xnatId <MR-session ID>'
    print '    -host <XNAT base URI>'
    print '    -u <XNAT username>'
    print '    -pwd <XNAT password>'
    print '    -raw <create QC files for raw scans only> (ignored)'
    print
    print 'new options are:'
    print
    print '    -scan <comma-separated list of scans> (may occur more than once)'
    print '    -nooverwrite'
    print '    -test <directory>'
    print '    -clear'
    print
    print 'existence and overwriting is done on a per-SNAPSHOTS resource basis '
    print '(so -nooverwrite will skip a scan if the resource exists, and '
    print 'overwriting will replace the whole resource)'
    print
    print '-clear will create a gray image (useful for testing)'
    print
    sys.exit(1)

session_id = None
session_label = None
base_uri = None
username = None
password = None
overwrite_flag = True
target_scans = None
test_dir = None

while sys.argv:
    option = sys.argv.pop(0)
    try:
        if option == '-xnatId':
            session_id = sys.argv.pop(0)
        elif option == '-session':
            session_label = sys.argv.pop(0)
        elif option == '-project':
            sys.argv.pop(0)
        elif option == '-host':
            base_uri = sys.argv.pop(0)
        elif option == '-u':
            username = sys.argv.pop(0)
        elif option == '-pwd':
            password = sys.argv.pop(0)
        elif option == '-raw':
            pass
        elif option == '-nooverwrite':
            overwrite_flag = False
        elif option == '-clear':
            generators = (gen_clear, )
        elif option == '-scan':
            if target_scans is None:
                target_scans = set()
            target_scans.update(sys.argv.pop(0).split(','))
        elif option == '-test':
            test_dir = sys.argv.pop(0)
        else:
            command_line_error('unknown option "%s"' % option)
            sys.exit(1)
    except IndexError:
        command_line_error('missing argument to %s' % option)
        sys.exit(1)

if test_dir is not None:
    print 'entering test mode'
    session_label = 'test-session'
    scans = []
    for subdir in os.listdir(test_dir):
        scans.append(DirScan(session_label, '%s/%s' % (test_dir, subdir)))
else:
    if not session_id:
        command_line_error('no session ID given')
        sys.exit(1)
    if not session_label:
        command_line_error('no session label given')
        sys.exit(1)
    if not base_uri:
        command_line_error('no XNAT URI given')
        sys.exit(1)
    if not username:
        command_line_error('no username given')
        sys.exit(1)
    if not password:
        command_line_error('no password given')
        sys.exit(1)
    connection = xnat.Connection(base_uri, username, password)
    e = connection.find_experiment(session_id)
    scans = [ XNATScan(session_label, s) for s in e.scans.itervalues() ]

for scan in scans:
    if target_scans is not None:
        if scan.id not in target_scans:
            print '*', scan, '...skipping'
            continue
        target_scans.remove(scan.id)
    print '*', scan
    if not overwrite_flag:
        if 'SNAPSHOTS' in scan.resources:
            print 'SNAPSHOTS exists', scan
            continue
    im = None
    try:
        for f in generators:
            try:
                im = f(scan)
            except GeneratorError, e:
                print '%s(): %s' % (f.__name__, e.reason)
            except StandardError:
                print '%s(): error' % f.__name__
                traceback.print_exc()
            else:
                if not im:
                    print '%s(): failed silently' % f.__name__
            if im:
                print '%s(): success' % f.__name__
                break
        if not im:
            print 'no generator happy with', scan
        else:
            scan.save_image(im)
    finally:
        scan.remove_copies()

if target_scans:
    print 'the following requested scans were not found:'
    for s in target_scans:
        print '    %s' % s

if test_dir is None:
    connection.close()

sys.exit(0)

# eof