Updated register script, added pixel size for Omero import

Author: folterj

biomero_importer/utils/importer.py

@@ -489,13 +489,15 @@ def import_to_omero(self, conn, file_path, target_id, target_type, uuid, transfe
     @connection
     def import_zarr(self, conn, uri, target, target_by_name=None, endpoint=None, nosignrequest=False):
         # Using https://github.com/BioNGFF/omero-import-utils/blob/main/metadata/register.py
-        from .register import load_attrs, register_image, register_plate, link_to_target, validate_endpoint
+        from .register import (load_attrs, register_image, register_plate, link_to_target, validate_endpoint,
+                               get_omexml_bytes, full_import, parse_image_metadata, set_rendering_settings,
+                               set_external_info)
         import zarr
         from types import SimpleNamespace
 
         file_title = os.path.splitext(os.path.basename(uri))[0].rstrip('.ome')
         args = SimpleNamespace(uri=uri, endpoint=endpoint, name=file_title,
-                               nosignrequest=nosignrequest, target=target, target_by_name=target_by_name)
+                               nosignrequest=nosignrequest, target=str(target), target_by_name=target_by_name)
 
         # --- start copy from register.main() ---
 
@@ -524,16 +526,43 @@ def import_zarr(self, conn, uri, target, target_by_name=None, endpoint=None, nos
         else:
             if "bioformats2raw.layout" in zattrs and zattrs["bioformats2raw.layout"] == 3:
                 print("Registering: bioformats2raw.layout")
-                series = 0
-                series_exists = True
-                while series_exists:
-                    try:
-                        print("Checking for series:", series)
-                        obj = register_image(conn, store, args, None, image_path=str(series))
-                        objs.append(obj)
-                    except FileNotFoundError:
-                        series_exists = False
-                    series += 1
+                zarr_name = args.uri.rstrip("/").split("/")[-1]
+                if args.name:
+                    zarr_name = args.name
+                # try to load OME/METADATA.ome.xml
+                omexml_bytes = get_omexml_bytes(store)
+                if omexml_bytes is not None:
+                    print("Importing OME/METADATA.ome.xml")
+                    rsp = full_import(conn.c, omexml_bytes, args.wait)
+                    for series, p in enumerate(rsp.pixels):
+                        # set external info. NB: order of pixels MUST match the series 0, 1, 2...
+                        image = conn.getObject("Image", p.image.id.val)
+                        image_path = str(series)
+                        image_attrs = load_attrs(store, image_path)
+                        # pixels_type is only used if we have *incomplete* `omero` metadata
+                        sizes, pixel_size, pixels_type = parse_image_metadata(store, image_attrs, image_path)
+                        rnd_def = set_rendering_settings(conn, image, image_attrs, pixels_type)
+                        if rnd_def is not None:
+                            conn.getUpdateService().saveAndReturnObject(rnd_def)
+                        set_external_info(image._obj, args, image_path=image_path)
+                        # default name is METADATA.ome.xml [series], based on clientPath?
+                        new_name = image.name.replace("METADATA.ome.xml", zarr_name)
+                        print("Imported Image:", image.id, new_name)
+                        image.setName(new_name)
+                        image.save()  # save Name and ExternalInfo
+                        objs.append(image)
+                else:
+                    print("OME/METADATA.ome.xml Not Found")
+                    series = 0
+                    series_exists = True
+                    while series_exists:
+                        try:
+                            print("Checking for series:", series)
+                            obj = register_image(conn, store, args, None, image_path=str(series))
+                            objs.append(obj)
+                        except FileNotFoundError:
+                            series_exists = False
+                        series += 1
             else:
                 print("Registering: Image")
                 objs = [register_image(conn, store, args, zattrs)]

biomero_importer/utils/register.py

@@ -21,16 +21,21 @@
 
 from urllib.parse import urlsplit
 import zarr
+from zarr.core.buffer import default_buffer_prototype
+from zarr.core.sync import sync
 
 import argparse
+import locale
+import platform
+import sys
+import tempfile
 
 from numpy import iinfo, finfo
 
 # from getpass import getpass
 
 from omero.cli import cli_login
 from omero.gateway import BlitzGateway
-from omero.gateway import OMERO_NUMPY_TYPES
 
 import omero
 
@@ -41,6 +46,12 @@
 
 from omero.model import ExternalInfoI
 from omero.rtypes import rbool, rdouble, rint, rlong, rstring, rtime
+from omero.model import ChecksumAlgorithmI
+from omero.model import LengthI
+from omero.model import NamedValue
+from omero.model.enums import ChecksumAlgorithmSHA1160
+from omero_version import omero_version
+from omero.callbacks import CmdCallbackI
 
 AWS_DEFAULT_ENDPOINT = "s3.us-east-1.amazonaws.com"
 
@@ -58,6 +69,123 @@
                'complex_': PixelsTypecomplex,
                'complex64': PixelsTypecomplex}
 
+
+def get_omexml_bytes(store):
+    # get() is async. Need to sync to get the bytes
+    rsp = store.get("OME/METADATA.ome.xml", prototype=default_buffer_prototype())
+    result = sync(rsp)
+    if result is None:
+        return None
+    return result.to_bytes()
+
+
+def create_fileset():
+    """Create a new Fileset with single OME XML file."""
+    fileset = omero.model.FilesetI()
+    entry = omero.model.FilesetEntryI()
+    # NB: If the clientPath includes .zarr, Bio-Formats tries to import zarr group
+    entry.setClientPath(rstring("OME/METADATA.ome.xml"))
+    fileset.addFilesetEntry(entry)
+
+    # Fill version info
+    system, node, release, version, machine, processor = platform.uname()
+
+    client_version_info = [
+        NamedValue('omero.version', omero_version),
+        NamedValue('os.name', system),
+        NamedValue('os.version', release),
+        NamedValue('os.architecture', machine)
+    ]
+    try:
+        client_version_info.append(
+            NamedValue('locale', locale.getlocale()[0]))
+    except:
+        pass
+
+    upload = omero.model.UploadJobI()
+    upload.setVersionInfo(client_version_info)
+    fileset.linkJob(upload)
+    return fileset
+
+
+def create_settings():
+    """Create ImportSettings and set some values."""
+    settings = omero.grid.ImportSettings()
+    # can't create thumbnails on import since ExternalInfo is not set yet
+    settings.doThumbnails = rbool(False)
+    settings.noStatsInfo = rbool(False)
+    settings.userSpecifiedTarget = None
+    settings.userSpecifiedName = None
+    settings.userSpecifiedDescription = None
+    settings.userSpecifiedAnnotationList = None
+    settings.userSpecifiedPixels = None
+    settings.checksumAlgorithm = ChecksumAlgorithmI()
+    s = rstring(ChecksumAlgorithmSHA1160)
+    settings.checksumAlgorithm.value = s
+    return settings
+
+
+def upload_file(proc, omexml_bytes, client):
+    """Upload files to OMERO from local filesystem."""
+    ret_val = []
+    i = 0
+    rfs = proc.getUploader(i)
+    try:
+        offset = 0
+        # rfs.write([], offset, 0)  # Touch
+        # Write the OME XML file
+        rfs.write(omexml_bytes, offset, len(omexml_bytes))
+
+        # create temp file for sha1
+        with tempfile.NamedTemporaryFile(delete_on_close=False) as fp:
+            fp.write(omexml_bytes)
+            fp.close()
+            ret_val.append(client.sha1(fp.name))
+    finally:
+        rfs.close()
+    return ret_val
+
+
+def assert_import(client, proc, omexml_bytes, wait):
+    """Wait and check that we imported an image."""
+    hashes = upload_file(proc, omexml_bytes, client)
+    # print ('Hashes:\n  %s' % '\n  '.join(hashes))
+    handle = proc.verifyUpload(hashes)
+    cb = CmdCallbackI(client, handle)
+
+    # https://github.com/openmicroscopy/openmicroscopy/blob/v5.4.9/components/blitz/src/ome/formats/importer/ImportLibrary.java#L631
+    if wait == 0:
+        cb.close(False)
+        return None
+    if wait < 0:
+        while not cb.block(2000):
+            sys.stdout.write('.')
+            sys.stdout.flush()
+        sys.stdout.write('\n')
+    else:
+        cb.loop(wait, 1000)
+    rsp = cb.getResponse()
+    if isinstance(rsp, omero.cmd.ERR):
+        raise Exception(rsp)
+    assert len(rsp.pixels) > 0
+    return rsp
+
+
+def full_import(client, omexml_bytes, wait=-1):
+    """Re-usable method for a basic import."""
+    mrepo = client.getManagedRepository()
+
+    fileset = create_fileset()
+    settings = create_settings()
+
+    proc = mrepo.importFileset(fileset, settings)
+    try:
+        # do the upload and trigger the import
+        return assert_import(client, proc, omexml_bytes, wait)
+    finally:
+        proc.close()
+
+
 def format_s3_uri(uri, endpoint):
     '''
     Combine endpoint and uri
@@ -109,8 +237,17 @@ def parse_image_metadata(store, img_attrs, image_path=None):
             else:
                 sizes[axis["name"]] = size
 
+    pixel_size = {}
+    transforms = multiscale_attrs["datasets"][0]["coordinateTransformations"]
+    for transform in transforms:
+        if transform["type"] == "scale":
+            scale = transform["scale"]
+            pixel_size = {axis["name"]: (pixel_size, axis.get("unit", "")) for axis, pixel_size
+                          in zip(axes, scale) if axis["name"] in "xyz"}
+            break
+
     pixels_type = array_data.dtype.name
-    return sizes, pixels_type
+    return sizes, pixel_size, pixels_type
 
 
 def create_image(conn, store, image_attrs, object_name, families, models, args, image_path=None):
@@ -119,7 +256,7 @@ def create_image(conn, store, image_attrs, object_name, families, models, args,
     '''
     query_service = conn.getQueryService()
     pixels_service = conn.getPixelsService()
-    sizes, pixels_type = parse_image_metadata(store, image_attrs, image_path)
+    sizes, pixel_size, pixels_type = parse_image_metadata(store, image_attrs, image_path)
     size_t = sizes.get("t", 1)
     size_z = sizes.get("z", 1)
     size_x = sizes.get("x", 1)
@@ -133,11 +270,10 @@ def create_image(conn, store, image_attrs, object_name, families, models, args,
 
     rnd_def = None
     image = conn.getObject("Image", iid)
-    omero_attrs = image_attrs.get('omero', None)
-    if omero_attrs is not None:
-        set_channel_names(conn, iid, omero_attrs)
-        # Check rendering settings
-        rnd_def = set_rendering_settings(omero_attrs, pixels_type, image.getPixelsId(), families, models)
+    # Set rendering settings and channel names if omero_attrs is provided
+    rnd_def = set_rendering_settings(conn, image, image_attrs, pixels_type, families, models)
+
+    set_pixel_size(image, pixel_size)
 
     img_obj = image._obj
     set_external_info(img_obj, args, image_path)
@@ -174,10 +310,22 @@ def set_channel_names(conn, iid, omero_attrs):
     conn.setChannelNames("Image", [iid], nameDict)
 
 
-def set_rendering_settings(omero_info, pixels_type, pixels_id, families, models):
+def set_rendering_settings(conn, image, image_attrs, pixels_type, families=None, models=None):
     '''
     Extract the rendering settings and the channels information
     '''
+    omero_info = image_attrs.get('omero', None)
+    if omero_info is None:
+        return None
+    set_channel_names(conn, image.id, omero_info)
+
+    if families is None:
+        families = load_families(conn.getQueryService())
+    if models is None:
+        models = load_models(conn.getQueryService())
+
+    pixels_id = image.getPrimaryPixels().getId()
+    
     if omero_info is None:
         return
     rdefs = omero_info.get('rdefs', None)
@@ -255,6 +403,16 @@ def set_rendering_settings(omero_info, pixels_type, pixels_id, families, models)
     return rnd_def
 
 
+def set_pixel_size(image, pixel_size):
+    pixels = image.getPrimaryPixels()._obj
+    if "x" in pixel_size:
+        pixels.setPhysicalSizeX(LengthI(pixel_size["x"][0], pixel_size["x"][1].upper()))
+    if "y" in pixel_size:
+        pixels.setPhysicalSizeY(LengthI(pixel_size["y"][0], pixel_size["y"][1].upper()))
+    if "z" in pixel_size:
+        pixels.setPhysicalSizeZ(LengthI(pixel_size["z"][0], pixel_size["z"][1].upper()))
+
+
 def load_families(query_service):
     ctx = {'omero.group': '-1'}
     return query_service.findAllByQuery('select f from Family as f', None, ctx)
@@ -498,9 +656,15 @@ def link_to_target(args, conn, obj):
 
     if args.target:
         if is_plate:
-            target = conn.getObject("Screen", attributes={'id': int(args.target)})
+            screen_id = args.target
+            if screen_id.startswith("Screen:"):
+                screen_id = screen_id.split(":")[1]
+            target = conn.getObject("Screen", attributes={'id': int(screen_id)})
         else:
-            target = conn.getObject("Dataset", attributes={'id': int(args.target)})
+            dataset_id = args.target
+            if dataset_id.startswith("Dataset:"):
+                dataset_id = dataset_id.split(":")[1]
+            target = conn.getObject("Dataset", attributes={'id': int(dataset_id)})
     else:
         if is_plate:
             target = conn.getObject("Screen", attributes={'name': args.target_by_name})
@@ -532,6 +696,10 @@ def main():
     parser.add_argument("--nosignrequest", required=False, action='store_true', help="Indicate to sign anonymously")
     parser.add_argument("--target", required=False, type=str, help="The id of the target (dataset/screen)")
     parser.add_argument("--target-by-name", required=False, type=str, help="The name of the target (dataset/screen)")
+    parser.add_argument('--wait', type=int, default=-1, help=(
+        'Wait for this number of seconds for each import to complete. '
+        '0: return immediately, -1: wait indefinitely (default). '
+        'Only applies when importing OME/METADATA.ome.xml.'))
 
     args = parser.parse_args()
 
@@ -565,16 +733,43 @@ def main():
         else:
             if "bioformats2raw.layout" in zattrs and zattrs["bioformats2raw.layout"] == 3:
                 print("Registering: bioformats2raw.layout")
-                series = 0
-                series_exists = True
-                while series_exists:
-                    try:
-                        print("Checking for series:", series)
-                        obj = register_image(conn, store, args, None, image_path=str(series))
-                        objs.append(obj)
-                    except FileNotFoundError:
-                        series_exists = False
-                    series += 1
+                zarr_name = args.uri.rstrip("/").split("/")[-1]
+                if args.name:
+                    zarr_name = args.name
+                # try to load OME/METADATA.ome.xml
+                omexml_bytes = get_omexml_bytes(store)
+                if omexml_bytes is not None:
+                    print("Importing OME/METADATA.ome.xml")
+                    rsp = full_import(conn.c, omexml_bytes, args.wait)
+                    for series, p in enumerate(rsp.pixels):
+                        # set external info. NB: order of pixels MUST match the series 0, 1, 2...
+                        image = conn.getObject("Image", p.image.id.val)
+                        image_path = str(series)
+                        image_attrs = load_attrs(store, image_path)
+                        # pixels_type is only used if we have *incomplete* `omero` metadata
+                        sizes, pixel_size, pixels_type = parse_image_metadata(store, image_attrs, image_path)
+                        rnd_def = set_rendering_settings(conn, image, image_attrs, pixels_type)
+                        if rnd_def is not None:
+                            conn.getUpdateService().saveAndReturnObject(rnd_def)
+                        set_external_info(image._obj, args, image_path=image_path)
+                        # default name is METADATA.ome.xml [series], based on clientPath?
+                        new_name = image.name.replace("METADATA.ome.xml", zarr_name)
+                        print("Imported Image:", image.id, new_name)
+                        image.setName(new_name)
+                        image.save()   # save Name and ExternalInfo
+                        objs.append(image)
+                else:
+                    print("OME/METADATA.ome.xml Not Found")
+                    series = 0
+                    series_exists = True
+                    while series_exists:
+                        try:
+                            print("Checking for series:", series)
+                            obj = register_image(conn, store, args, None, image_path=str(series))
+                            objs.append(obj)
+                        except FileNotFoundError:
+                            series_exists = False
+                        series += 1
             else:
                 print("Registering: Image")
                 objs = [register_image(conn, store, args, zattrs)]