OMERO mask utility

mask_utils.py

# -*- coding: utf-8 -*-
"""
-----------------------------------------------------------------------------
  Copyright (C) 2016 Glencoe Software, Inc. All rights reserved.
  This program is free software; you can redistribute it and/or modify
  it under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.
  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU General Public License for more details.
  You should have received a copy of the GNU General Public License along
  with this program; if not, write to the Free Software Foundation, Inc.,
  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
------------------------------------------------------------------------------
"""


import os
import argparse

from PIL import Image
import math
import struct

import omero
import omero.clients
from omero.gateway import BlitzGateway
from omero.gateway import ColorHolder
from omero.rtypes import rint, rdouble
import numpy


color_table = {}
color_table['red'] = ColorHolder()
color_table['blue'] = ColorHolder()
color_table['green'] = ColorHolder()

color_table['red'].setRed(255)
color_table['red'].setBlue(0)
color_table['red'].setGreen(255)
color_table['red'].setAlpha(100)

color_table['green'].setRed(0)
color_table['green'].setBlue(0)
color_table['green'].setGreen(255)
color_table['green'].setAlpha(100)

color_table['blue'].setRed(0)
color_table['blue'].setBlue(255)
color_table['blue'].setGreen(0)
color_table['blue'].setAlpha(100)


def create_mask(
        gateway, file_path, image_id, channel,
        z_plane, timepoint, color, threshold, x, y):
    '''
    Creates a mask from the image file and attaches it to the OMERO image.
    '''
    image = gateway.getObject("Image", image_id)
    pil_image = Image.open(file_path)
    width, height = pil_image.size
    print "Attaching mask size: %s, %s" % (width, height)
    if '16' in pil_image.mode:
        divider = 16.0
        format_string = "H"  # Unsigned short
        byte_factor = 0.5
    else:
        divider = 8.0
        format_string = "B"  # Unsiged char
        byte_factor = 1
    bytes = pil_image.tobytes()
    steps = math.ceil(len(bytes) / divider)
    mask_array = []
    for i in range(long(steps)):
        binary = bytes[i * int(divider):i * int(divider) + int(divider)]
        format = str(int(byte_factor * len(binary))) + format_string
        binary = struct.unpack(format, binary)
        s = ""
        for bit in binary:
            if bit >= threshold:
                bit = 1
            else:
                bit = 0
            s += str(bit)
        mask_array.append(int(s, 2))
    mask_bytes = bytearray(mask_array)
    roi = omero.model.RoiI()
    mask = omero.model.MaskI()
    mask.setTheC(rint(channel))
    mask.setTheZ(rint(z_plane))
    mask.setTheT(rint(timepoint))
    mask.setX(rdouble(x))
    mask.setY(rdouble(y))
    mask.setWidth(rdouble(width))
    mask.setHeight(rdouble(height))
    mask.setFillColor(rint(color_table[color].getInt()))
    mask.setBytes(mask_bytes)
    roi.addShape(mask)
    roi.setImage(image._obj)
    roi = gateway.getUpdateService().saveAndReturnObject(roi)
    print "Attach masks to image %s, new ROI id: %s" \
        % (image_id, roi.getId().getValue())


def write_mask(gateway, roi_id):
    roi_service = gateway.getRoiService()
    result = roi_service.findByRoi(roi_id, None)
    for roi in result.rois:
        for shape in roi.copyShapes():
            if type(shape) != omero.model.MaskI:
                continue
            width = int(shape.getWidth().getValue())
            height = int(shape.getHeight().getValue())
            bytes = shape.getBytes()
            byte_array = numpy.fromstring(bytes, dtype=numpy.uint8)
            binary = numpy.unpackbits(byte_array)
            binary = binary.reshape((width, height))
            im = Image.fromarray(255 * binary)
            mask_name = "mask_%s_%s.png" % (roi_id, shape.getId().getValue())
            im.save(mask_name)
            print "Saved mask id: %s to file: %s" \
                % (shape.getId().getValue(), mask_name)


def parse_arguments():
    description = "Upload/download masks to and from OMERO server.\n" \
        "Supports only 8bit and 16bit grayscale images."
    width = 225
    column_width = 45
    parser = argparse.ArgumentParser(
        prog='python mask_utils.py', description=description,
        formatter_class=lambda
        prog: argparse.HelpFormatter(
            prog, max_help_position=column_width, width=width
        )
    )
    parser.add_argument(
        '-s', '--server', required=True,
        help='OMERO server.'
    )
    parser.add_argument(
        '-p', '--port',
        help='OMERO port.'
    )
    parser.add_argument(
        '-u', '--user',
        help='OMERO user.'
    )
    parser.add_argument(
        '-w', '--password',
        help='OMERO password.'
    )
    parser.add_argument(
        '-k', '--key',
        help='OMERO session UUID.'
    )
    
    subparsers = parser.add_subparsers(dest='command')
    
    parser_upload = subparsers.add_parser(
        'upload_mask', help = 'Upload mask to OMERO.',
        formatter_class=lambda 
        prog: argparse.HelpFormatter(
            prog, max_help_position=column_width, width=width
        )
    )
    parser_upload.add_argument(
        '--file', '-f', required=True,
        help='File containing the mask.'
    )
    parser_upload.add_argument(
        '--image', '-i', required=True, type=int,
        help='OMERO image Id to attach the mask to.'
    )
    parser_upload.add_argument(
        '--channel', '-c', default=0, type=int,
        help='OMERO image channel to attach the mask to.'
    )
    parser_upload.add_argument(
        '--timepoint', '-t', default=0, type=int,
        help='OMERO image time-point to attach the mask to.'
    )
    parser_upload.add_argument(
        '--z_plane', '-z', default=0, type=int,
        help='OMERO z plane to attach the mask to.'
    )
    parser_upload.add_argument(
        '-x', default=0, type=int,
        help='X coordinate.'
    )
    parser_upload.add_argument(
        '-y', default=0, type=int,
        help='Y coordinate.'
    )
    parser_upload.add_argument(
        '--color', required=True, choices=['red','green', 'blue'],
        help="Mask's color"
    )
    parser_upload.add_argument(
        '--threshold', default=1, type=int,
        help='Threshold to apply to image when creating a mask.'
    )
    
    parser_download= subparsers.add_parser(
        'download_mask', help = 'Download mask from OMERO.',
        formatter_class=lambda
        prog: argparse.HelpFormatter(
            prog, max_help_position=column_width, width=width
        )
    )
    parser_download.add_argument(
        '--roi_id', '-id', required=True, type=int,
        help='ROI id containing the mask.'
    )
    args = parser.parse_args()
    return vars(args)


def setup_connection(params):
    conn = omero.client(host=params["server"], port=params["port"])
    session = None
    if params["key"] is not None:
        conn.joinSession(params["key"])
    elif params["user"] is not None \
            and params["password"] is not None:
        session = conn.createSession(
            username=params["user"],
            password=params["password"]
        )
        conn.enableKeepAlive(60)
        session.detachOnDestroy()
    gateway = BlitzGateway(client_obj=conn)
    user = gateway.getUser()
    print "Current user:"
    print "   ID:", user.getId()
    print "   Username:", user.getName()
    print "   Full Name:", user.getFullName()
    group = gateway.getGroupFromContext()
    print "Current group: ", group.getName()
    return gateway


def upload_mask(params):
    gateway = setup_connection(params)
    try:
        create_mask(
            gateway, params["file"], params["image"], params["channel"],
            params["timepoint"], params["z_plane"], params["color"],
            params["threshold"], params["x"], params["y"]
        )
    finally:
        gateway._closeSession()


def download_mask(params):
    gateway = setup_connection(params)
    try:
        write_mask(gateway, params["roi_id"])
    finally:
        gateway._closeSession()


if __name__ == "__main__":
    process_descriptor = parse_arguments()
    if process_descriptor["command"] == "upload_mask":
        upload_mask(process_descriptor)
    if process_descriptor["command"] == "download_mask":
        download_mask(process_descriptor)