From 9b70e2b8ee467383ddbe3c40163c80e9a1b4e110 Mon Sep 17 00:00:00 2001
From: davelopez <46503462+davelopez@users.noreply.github.com>
Date: Thu, 5 Feb 2026 16:39:17 +0100
Subject: [PATCH] Streamlines OME-TIFF export to reduce memory usage

Replaces full image array allocation with a generator-based streaming approach for exporting multi-dimensional OME-TIFF files. This minimizes memory consumption by writing planes one at a time.
---
 lib/galaxy/files/sources/omero.py | 38 +++++++++++++++++++------------
 1 file changed, 23 insertions(+), 15 deletions(-)

diff --git a/lib/galaxy/files/sources/omero.py b/lib/galaxy/files/sources/omero.py
index 6104cf6b5d3..38f9d486da6 100644
--- a/lib/galaxy/files/sources/omero.py
+++ b/lib/galaxy/files/sources/omero.py
@@ -6,7 +6,6 @@ from typing import (
     Union,
 )
 
-import numpy as np
 import tifffile
 
 from galaxy.exceptions import (
@@ -524,14 +523,13 @@ class OmeroFileSource(BaseFilesSource[OmeroFileSourceTemplateConfiguration, Omer
                 f.write(chunk)
 
     def _export_as_tiff(self, image, native_path: str):
-        """Export all Z-planes and channels of the image as a multi-dimensional TIFF.
+        """Export all Z-planes and channels of the image as a multi-dimensional OME-TIFF.
 
         Exports all channels (C) and Z-planes at the first timepoint (T=0).
-        The image is saved with ZCYX axis ordering, which is the standard format
-        expected by Galaxy's image handling tools.
+        Uses generator-based streaming to minimize memory usage - the full shape
+        is declared upfront but planes are fetched and written one at a time.
 
-        This method preserves multichannel information that would otherwise be lost
-        when downloading from servers like IDR that restrict original file access.
+        Memory usage: O(Y*X) per plane instead of O(Z*C*Y*X) for entire image.
         """
         pixels = image.getPrimaryPixels()
         size_z = image.getSizeZ()
@@ -539,17 +537,27 @@ class OmeroFileSource(BaseFilesSource[OmeroFileSourceTemplateConfiguration, Omer
         size_y = image.getSizeY()
         size_x = image.getSizeX()
 
-        # Create array with shape (Z, C, Y, X) to hold all planes
-        image_array = np.zeros((size_z, size_c, size_y, size_x), dtype=pixels.getPlane(0, 0, 0).dtype)
+        # Get dtype from first plane sample
+        first_plane = pixels.getPlane(0, 0, 0)
+        dtype = first_plane.dtype
 
-        # Fetch all Z-planes and channels from OMERO
-        for z in range(size_z):
-            for c in range(size_c):
-                plane_data = pixels.getPlane(z, c, 0)  # z, channel, timepoint
-                image_array[z, c, :, :] = plane_data
+        def plane_generator():
+            """Generator yielding 2D planes in ZCYX order (Z outer, C inner)."""
+            for z in range(size_z):
+                for c in range(size_c):
+                    if z == 0 and c == 0:
+                        # First plane already fetched for dtype detection
+                        yield first_plane
+                    else:
+                        yield pixels.getPlane(z, c, 0)
 
-        # Write as TIFF with proper axis metadata
-        tifffile.imwrite(native_path, image_array, metadata={"axes": "ZCYX"})
+        with tifffile.TiffWriter(native_path, bigtiff=True, ome=True) as tif:
+            tif.write(
+                data=plane_generator(),
+                shape=(size_z, size_c, size_y, size_x),
+                dtype=dtype,
+                metadata={"axes": "ZCYX"},
+            )
 
     def _write_from(
         self, target_path: str, native_path: str, context: FilesSourceRuntimeContext[OmeroFileSourceConfiguration]