Illumination correction task

[tcompa]

@gusqgm already explained the basics of the illumination correction task in #26 (comment). Let's see if we are all on the same page.
In our understanding, the steps are as follows:

1. Load correction matrix (which has shape YX: 2160x2560), and normalize it to go from its initial range (e.g. [138,255]) to the range [0,1].
2. Within the zarr array (created by the yokogawa_to_zarr task), subtract the background (=120, as a default value) from all the array elements.
3. Within the zarr array, and for each 2160x2560 region, divide the background-subtracted image by the normalized correction matrix. Notice that array regions could correspond to array chunks, but this is not necessarily true (see https://github.com/fractal-analytics-platform/mwe_fractal/issues/59).

Concerning the output, our first proposal is that we just replace the original zarr array (the one obtained via yokogawa_to_zarr) with the new corrected one. Later on we can decide how to proceed: does it have to be a on/off flag provided by the user (see also https://github.com/fractal-analytics-platform/mwe_fractal/issues/29)? Note that this will require some care, since (by now) creating a new zarr file requires a global task that first generates its structure (as we currently do with replicate_zarr_structure_mip for the MIP task).

[gusqgm]

Hey @tcompa , thanks for the clear step description!

This is currently what we do at the Liberali lab indeed.
However we have been looking at other ways on how to correct for the microscope images as well, as our current method might be not the end solution. This will not affect input, output, only minor changes on how the correction matrix should be applied onto the acquired images, so should not change the discussion here conceptually.

Concerning the more important point of the output strategy: since ideally the illumination correction should be part of a larger "PrepProcessing Workflow" for the lab (at least this is what happens not at the Liberalis), rewriting the .zarr structure should be fine for now, so the above strategy is sound!

The flexibility of writing the output onto a new .zarr file is still important though as I can imagine that it can: 1) aid us in debugging once we have more tasks added into the task list and 2) aid the user in general in learning about optimal processing parameters when using a novel task / learning things from scratch. @tcompa , how difficult do you think it would be to create a bundled task that replicates and outputs the new .zarr structure considering the processed data? Also, do you think it would make sense to have a pointer file that would refer to these different .zarr structures (i.e. in different processing states) so that a user could best refer to which files to further process / delete? Would the naming of the folder be enough?

Any other thoughts?

[tcompa]

Thanks for these details.

I agree: storing the corrected data in a new zarr file has relevant use cases. But it also opens a broader discussion that we are having
Let's touch on a few of its aspects:

1. In terms of tasks, there's nothing complex in what you suggest, i.e. having two different tasks, one that recreates the new zarr structure (parallelized at the plate level) and one that performs the illumination correction (parallelized at the well level) and writes output in the new array. Bundling them together simply requires executing them one after the other.
2. Having a simple on/off switch that decides whether the output should go to the new zarr (which requires creating it) or to the old one (which requires no additional step) is trivial. We add a True/False parameter, and in one case we add something like replicate_zarr_structure_illumination before the illumination_correction task.
3. In terms of workflow, this is a change with respect to our current infrastructure. What we are doing at the moment is: tasks = python functions, workflow = list of tasks (with their arguments and properties) The only exception (which we plan to remove at some later point) is that fractal_cmd.py also handles the additional parameters for each task, in a task-specific way. If we add the on/off switch that possibly introduces a new task in the workflow (the one for generating new zarr structures, but only if needed), we are actually moving part of the task logic into fractal_cmd.py (or its future versions), which now goes beyond the simple goal of encoding/validating/executing a workflow. No problem with that (and it's likely that it is inevitable, especially if we want things like skipping the execution of tasks that already ran), but let's define clearly how much of it we are happy with.

Note that this also applies to the MIP task (first a global task to replicate the zarr, then one for the actual MIP), but in that case it seems less likely that one wants to add the 2D data in the same zarr array as the 3D ones.

Something which is different, but also very important, is the issue of the pointer files. We can imagine having a few zarr files for each plate (e.g. the raw one, possibly the corrected one, one for the MIP, .. what else?), and tasks should "know" which one to use. At the moment this is all in the folder names (some are .zarr, some are _mip.zarr), but it needs more thinking.

Some remarks:

• This is worth a bit of discussion in our weekly meeting.
• We want to give the highest priority to expanding the available tasks, for a little while. Having this kind of discussions on the infrastructure is crucial, and it has to be done, but this may not lead to immediate change to the fractal_cmd.py code.

[tcompa]

A couple of quick questions on the output (showing my little experience with image analysis..):

1. Dividing an image (AKA an array of uint16 numbers) by a correction matrix in the [0,1] range leads to a set of floats over a large range. Is there any reasonable normalization to be applied to the output, or is it just (image-background)/correction?
2. Is there any reason to convert the output back to some integer format (apart from a huge gain in disk space), or should it remain an array of float numbers?
3. The (normalized) correction matrix takes values starting from 0, meaning that in principle there will be some 1/0 divergence in the correction. How should we treat this? Should we set a cutoff on small values of the matrix? (something like setting the normalized matrix equal to 1 for all values below 1e-10, or whatever else).

Well, upon writing this I realize that here I am normalizing the correction matrix range from its original one (say [138,255]) to [0,1], meaning that 138 is mapped to 0. But this is probably not what @gusqgm meant, and I guess I should rather divide by the maximum value (without subtracting 138). This seems much more reasonable.. Gustavo, could you please confirm?

[tcompa]

To be more explicit, I was doing:

        normalized_corr = (corr - corr.min()) / (corr.max() - corr.min())

but I guess the correct thing to do is

        normalized_corr = corr / corr.max()

[gusqgm]

Hey @tcompa you are right, dividing by the maximum yields a better result, especially considering that at this stage you have already performed the sole subtraction of the camera background values, so forcing more values to 0 while normalizing will most likely lead to loss of information in some cases.

Regarding some of the other above points:

2. usually the corrected output is rescaled back to the original bit depth, which in our case is then 16bit. It is worth to note that conceptually nothing prevents to go to an even smaller bit depth (e.g. 12bit) as long as we do not lose information there, but I do not see a need to try to optimize this now.

3. There should be a cut-off in my opinion, and I will get more information about it tomorrow. Also, let's keep in mind that there are other ways to perform background subtraction and currently I am checking the efficacy on others as well.

[tcompa]

Thanks. At the moment, my understanding is the followingç

• The maximum correction is approximately a x2 factor, because the correction matrices have ranges like [138,255].
• The original images are made of uint16 unsigned integers (between 0 and 65535).
• After correction, images are made of floats which are (roughly) in the range [0,2*65535], at least with the correction matrices I am using.

I have doubts on how to rescale the corrected images back to 16bit.

1. Should I simply do img = img * (65535 / img.max()) and then cast it as a uint16?
2. Should this conversion take place at the image level or at the level of an entire well/plate? That is, should there be img.max() in the rescaling factor or rather the maximum over the whole well/plate array.

[jluethi]

Concerning the implementation of the core illumination correction:
This page is a nice overview: https://clouard.users.greyc.fr/Pantheon/experiments/illumination-correction/index-en.html

The 1. Correction from a Dark Image and a Bright Image would be interesting if we had more inputs (e.g. a dark image as well as a bright image for control). What we have now is just the bright image, so we will do 2. Correction from a Bright Image (also called flat field correction).

Basically, it means:

I quickly tested this locally and it produces an output as we want it (it's not perfect for my data, because it was acquired on a different microscope, but you can see how it compensates at the edge). :
Pre correction:

Post correction:

Code for correction:

# Load images somehow
img_path = '/Users/joel/shares/dataShareFractal/2D/hiPSC_Slice/Cycle1/20200812-CardiomyocyteDifferentiation14-Cycle1_B03_T0001F001L01A01Z06C01.png'
img = imageio.v2.imread(img_path)
illum_corr_path = '/Users/joel/Desktop/IllumCorr/220131_40x_BP600_CH03.tif'
illum_img = imageio.v2.imread(illum_corr_path)

# Actual illumination correction:
img_corr = img/illum_img * np.mean(img) * 1/np.mean(img/illum_img)

# Recasting into uint16 (mostly stays in uint16, but in some edge cases it doesn't and recasting is useful, see next post)
img_corr = img_corr.astype('uint16')

The major benefits from this solution:

1. The values of the correct image are very similar to the input values
2. No need for clipping an image or rescaling it further
3. No need to [0, 1] normalize the illum_img before

[jluethi]

No, on the discussion of background subtraction: Yes, the illumination correction tends to work better after background subtraction. But users may vary in what they want for background subtraction. 120 definitely is too high of a default, I'd suggest we go with default background subtraction = 110. But it should be a parameter the user can change and 0 should be a valid option (= no background subtraction).

When doing background subtraction with uint16 images, one easily gets overflows (e.g. 100- 110 is 65325 in uint16. That's both wrong, looks weird and leads to the image no longer being in uint16 after illumination correction). Thus, I do background subtraction like this:

bg_threshold = 110
# Background subtraction: Careful about overflows! img - 120 leads to many values in range 65000!
img[img <= bg_threshold] = 0
img[img > bg_threshold] = img[img > bg_threshold] - 110

The illumination correction then looks a bit better using the method above:

[jluethi]

=> Let's go with this illumination correction & background subtraction approach for now, it seems to work quite well for general use-cases. We can implement more complicated versions later, e.g. calculating statistics on images (see here: https://pelkmanslab.org/wp-content/uploads/2019/02/StoegerBattichYakimovich2015.pdf) or use 1. Correction from a Dark Image and a Bright Image (as described here: https://clouard.users.greyc.fr/Pantheon/experiments/illumination-correction/index-en.html)

[tcompa]

Quick update:

The MWE for the illumination-correction task is essentially complete as it is in main, up to small improvements not yet on git. However we hit a subtle problem which is still not solved, and that's why the task is not ready to be tested at scale.

The issue appears when overwriting a zarr file (overwrite=True, in https://docs.dask.org/en/stable/generated/dask.array.to_zarr.html) with dask arrays.

• If we use to_zarr with a dask array for which we already forced execution (with an explicit .compute(), somewhere in the middle of the processing), then the overwriting works fine, but this would not be a good strategy for large arrays because it needs to load an entire well into memory.
• If we don't use any .compute() (which is in principle the correct way to deal with dask arrays), then the final to_zarr(..., overwrite=True) writes an array which has the correct shape but essentially empty chunks (that is, each chunk file is 44K instead of a few M).

Debugging this tricky issue is taking us more than expected, but it is a core feature that we definitely need to sort out ASAP, as it will likely appear within several other tasks.

PS
This is the issue I quickly mentioned in #36 (comment) and following discussion.

[tcompa]

Some more details.

What we are trying to do is essentially this:

import dask.array as da

# Save an array of zeros in a zarr file
x = da.zeros((100, 100), chunks=(10, 10), dtype=int)
x.to_zarr("main.zarr", component="level0", dimension_separator="/")

# Create a new array
z = da.ones((100, 100), chunks=(10, 10), dtype=int) * 2 + x 

# Overwrite previous zarr component
z.to_zarr("main.zarr", component="level0", dimension_separator="/", overwrite=True)

# Check that ovewriting was successful
new_array = da.from_zarr("main.zarr/level0")
assert new_array[0, 0].compute() == 2

This simple example does work, but the same thing fails when the new array comes from the more complex procedure we are using (involving several stacks, and a map_blocks to apply the correction to all images).

More on this soon, hopefully.

[tcompa]

Here is a much simpler minimal not-working example:

import shutil
import numpy as np
import dask.array as da

# Save an array to a zarr file
raw_data = da.random.random((3, 10, 400, 400), chunks=(1, 1, 200, 200))
raw_data = (raw_data * 10000).astype("uint16")
raw_data.to_zarr("raw.zarr", dimension_separator="/")

# Copy raw.zarr to main.zarr
shutil.copytree("raw.zarr", "main.zarr")

# Store the same data (with overwriting)
da.from_zarr("main.zarr").to_zarr("main.zarr", dimension_separator="/", overwrite=True)

After this script is run, we expect main.zarr and raw.zarr to have the same size (~10M), but we observe that main.zarr (the overwriten one) is much smaller (<1M). Note that if do to_zarr for another kind of array (i.e. not one loaded with from_zarr) things work as expected.

[tcompa]

Very relevant: dask/dask#5942

[tcompa]

The dask issue I mentioned (dask/dask#5942) is still open, after two years. It can be solved by patching dask/array/core.py locally (as in dask/dask@b9ef5be), but that seems a very risky way to go.

After a lot of time spent in debugging, at the moment me and @mfranzon opted for a quick&dirty custom solution for the overwrite=True case. Instead of actually overwriting, we write to a new component of the zarr array, then remove the original one, then rename the new one. See these lines in illumination_correction.py:

    # Write data into output zarr
    for ind_level in range(num_levels):
        if overwrite:
            pyramid[ind_level].to_zarr(
                newzarrurl,
                component=f"{ind_level}_TEMPORARY/",
                dimension_separator="/",
            )
            shutil.rmtree(newzarrurl + f"{ind_level}")
            shutil.move(
                newzarrurl + f"{ind_level}_TEMPORARY",
                newzarrurl + f"{ind_level}",
            )
        else:
            pyramid[ind_level].to_zarr(
                newzarrurl,
                component=f"{ind_level}",
                dimension_separator="/",
            )

It's clearly not the ideal way to go, but possibly sufficient for our needs. We are now testing it on the FMI dataset, to verify that at least the output looks reasonable.

[tcompa]

The example example_fmi_2_well_4x5_sites_illumination.sh (updated with f833409) performs the illumination correction (with overwrite) and then a MIP on the output. Its memory use is completely under control, even with wells of 30-40 G. Its CPU usage is still somewhat disappointing for a long time before spiking, but that's another issue. Here's the monitoring of the illumination correction on a single well:

I'd say that we are at the level where additional tests would be useful.

Important caveat: at the moment we are only loading FMI correction matrices for 60x and for channels 1,2,3,4 (each channel being corrected with the corresponding matrix). Next thing to do is to add an argument that somehow answers the question "which correction matrices should I use?". In fact there is also an example on git for the UZH dataset, but it's still using the FMI matrices and then there is no reason why it should produce nice output.

[tcompa]

Example of output (keep in mind that the two wells should be farther apart, in principle):

[jluethi]

That's great! We will probably need this ability to overwrite a zarr in different places, would be useful to abstract our workaround. Then, we can use this workaround with the temporary solution while there are still issues and, if the issues are fixed, we just replace our wrapper with the direct implementation :)

Also looks good regarding memory usage and the images look nice.
(Small side remark regarding the wells: If you add more columns to the .zattrs, it will display the wells with correct spacing. e.g. at the moment it only has columns 3 & 5. If we add 4, it then adds the gap. According to the OME-spec, all possible columns should be in that metadata, but whether this always makes sense is a different question. If we do that, take care that in the metadata, the column index of the second well also goes up by one in that case though, so maybe a slightly broader question to implement that)

Also, I will have a look whether we can get similarly structured illumination correction data for UZH images as well. We typically used a different approach, but maybe I can manually create some in the same format for the moment that fit the UZH data :)

[tcompa]

@jluethi, thanks for 20e36b6.

Quick answers:

1. How is the illum_corr matrix saved?
   I have to update this part. At the moment the path is just hardcoded, but it will be an input argument. In fact the input should be something like the initial part of the path, because then the end depends on the channels. I'll work out a simple example, and then we can make it more precise.
2. this is the same matrix everywhere for all channels!
   The corrections[int(chl)] part is taking care of the channel selection, isn't it? Anyway, I agree that it would be much nicer to define the correct function independently from the illumination_correction one, and that they should not share variables in any implicit way. I'll turn it into something like def correct(img, illum_img, (other parameters..)), and then illumination_correction should only take care of checking all shapes/chunking, identifying the correct channels and matrices, applying the correction, writing output.
3. Necessary to specify no Z?
   I opted for 3D images with a dummy Z index (image size along Z is of just one layer), because in this way I can use map_blocks naively (it takes a 3D array, and acts on its 3D chunks). In principle we can also explicitly loop over Z layers and apply map_blocks to 2D arrays/chunks, or try and find how to use it in a way that does not preserve shapes (that is, it takes a 2D image and spits out a 3D chunk.. or something like that). But the current problem (we need to see a 2D image as a 3D chunk with a dummy index) is very mild, so I don't see this is as an urgent issue.

[jluethi]

1. Yes, let's talk about this. I think a path to a folder is certainly necessary. But we will eventually have cases where more than 1 potential correction file is available. For example: We calculate illumination biases at the microscopes at different times over the year, sometimes also for a given experiment. Thus, there may be cases where multiple files are available and depending on when the experiment was acquired, a different correction file should be used. In other cases, users may acquire a correction file for their specific experiment or we may even add a task that can infer illumination profiles based on the existing images in the OME-Zarr file.

2. That was your comment, I just kept it in. But yes, those need to be channel specific. I actually had correct with illum_img as a parameter initially and removed it to fit the current scheme, but happy if that returns and the function is independent of the wrapping, that's a better separation of concerns. We can then also have illumination_correction handle parsing of which illum_img is loaded and the correction function just applies it.

3. I wasn't sure how you saved & read the illum_corr files, but this 3D approach seems to work well. No need to change that from my side :)

[tcompa]

Quick question (@gusqgm and @jluethi) about correction-matrix filenames.

Say that we have the files in the folders

/data/active/fractal/Liberali/FractalTesting20220124/IlluminationCorrectionMatrices-Yokogawa/

and filenames are like 220120_60xW_BP445_CH01.tif. The naive way to go is to provide an argument like initial_path to illumination_correction, which in this case would look like

initial_path = /data/active/fractal/Liberali/FractalTesting20220124/IlluminationCorrectionMatrices-Yokogawa/220120_60xW_BP445

We can then use glob to list all remaining files, which would differ only by the channel.

QUESTION
Is 220120_60xW_BP445_CH01.tif a standard filename that we can parse to extract the channel?
If there are different versions (e.g. for FMI and UZH microscopes), could you please make each one of them explicit (as in #48)?

[tcompa]

(the last comment partly overlaps with the simultaneous comment by Joel)

[jluethi]

Small overlap. Maybe the initial implementation is that there is just 1 illumination correction file per channel, but we will definitely need to support the use-case where more than 1 correction file exists. There may be a default way of handling it, but the user would need the ability to choose the file.

Regarding how standardized file naming is:
There are 3 conceptual ways to get to correction files, they may all have somewhat different filenames. But we can come up with a specification of how we want to name them and rename them on "import" to the system.

As an overview of where this profile may come from:

1. Illumination profile from microscope maintenance (e.g. the file we currently have, can be named 220120_60xW_BP445_CH01.tif). There can be multiple files per objective, laser & channel combination that vary in time they have been recorded.
2. Illumination profiles calculated based on specific hardware measured at the microscope (e.g. argolight). This is WIP, I don't think we have good examples of this so far, do we @gusqgm ?
3. Illumination profiles calculated by analyzing illumination statistics about many images in an experiment, e.g. measuring profiles of 10'000 images (often done at UZH at the moment) => generates an illumination correction file for the given experiment

Implementing 3 is not a priority for the moment and 2 is WIP at the FMI at the moment. So if we come up with a good way of defining how the files need to be named for 1 (e.g. similar to what we have there), we can then use that for approaches 2 & 3.
For approach 1, I'd say we need a path parameter and then a parameter for the full filename. The system could e.g. have a list of all available correction files sorted by date (first part of the filename) for a given objective (e.g. 60xW, 20x etc.), filterset (e.g. BP445) and channel (e.g. CH01). Maybe each microscope is in a different folder, so one would specify the folder for a given microscope. And when e.g. using an eventual approach 3, one would specify the output of that calculation as the illumination correction input folder

[tcompa]

As of our call today:
Let's aim for a minimal version where the user can provide a path for the correction-matrices folder and a set of files (one per channel).

[tcompa]

Let's aim for a minimal version where the user can provide a path for the correction-matrices folder and a set of files (one per channel).

Current version (536897e) does that.

At the moment there's only a quick&dirty way of passing inputs, which we can later improve on.

$ cat wf_params_uzh_1_well_2x2_sites.json
{
"workflow_name": "uzh_1_well_2x2_sites",
"dims": [2, 2],
"coarsening_xy": 3,
"coarsening_z": 1,
"num_levels": 5,
"channel_file": "wf_params_uzh_1_well_2x2_sites_channels.json",
"path_dict_corr": "wf_params_uzh_1_well_2x2_sites_illumination.json"
}

$ cat wf_params_uzh_1_well_2x2_sites_illumination.json 
{
    "root_path_corr": "/data/active/fractal/Liberali/FractalTesting20220124/IlluminationCorrectionMatrices-Yokogawa/",
    "A01_C01": "220120_60xW_BP445_CH01.tif",
    "A01_C02": "220120_60xW_BP525_CH02.tif",
    "A02_C03": "220120_60xW_BP600_CH03.tif"
}

and the actual task call is

def illumination_correction(
    zarrurl,
    overwrite=False,
    newzarrurl=None,
    chl_list=None,
    path_dict_corr=None,
    coarsening_xy=2,
    background=110,
):

    """
    Perform illumination correction of the array in zarrurl.

    :param zarrurl: input zarr, at the site level (e.g. x.zarr/B/03/0/)
    :type zarrurl: str
    :param overwrite: whether to overwrite an existing zarr file
    :type overwrite: bool
    :param newzarrurl: output zarr, at the site level (e.g. x.zarr/B/03/0/)
    :type newzarrurl: str
    :param chl_list: list of channel names (e.g. A01_C01)
    :type chl_list: list
    :param path_dict_corr: path of JSON file with info on illumination matrices
    :type path_dict_corr: str
    :param coarsening_xy: coarsening factor in XY (optional, default 2)
    :type coarsening_z: xy
    :param background: value for background subtraction (optional, default 110)
    :type background: int
    """

[tcompa]

Concerning the illumination-correction task, it would be useful to have another example of matrices, to see whether the current version (simple and very explicit in what user should provide) works reasonably.

@jluethi, perhaps we could have some illumination matrices related to UZH data?

After this test, I think we could close this issue for now, and we can reopen it later if we want to improve something. Comments?

[tcompa]

I'd like to test the FMI data with the correct illumination matrices. @gusqgm, do you confirm that this is the correct mapping between channels and files? Thanks

{
    "root_path_corr": "/data/active/fractal/Liberali/FractalTesting20220124/IlluminationCorrectionMatrices-Yokogawa/",
    "A04_C01": "220120_60xW_BP445_CH01.tif",
    "A03_C02": "220120_60xW_BP525_CH02.tif",
    "A02_C03": "220120_60xW_BP600_CH03.tif",
    "A01_C04": "220120_60xW_BP600_CH03.tif"
}

[jluethi]

@tcompa I created some mock illumination correction profiles for UZH (mock: not our real process to create them, but they are fitting the data very well)*. I updated all the UZH example scripts and created json files containing this info. Also used this opportunity to restructure the file naming a bit of the additional json files: Channels & illum_corr are not specific to a subset, but to the whole cardiac experiment. Thus, all cardiac test data now use those same json files.
Also, I updated the rescaling values (window start & end), because the appropriate values change with the background subtraction that we do during illumination correction.

As an illustration to show how well this works:

1. 2x2 sites without background subtraction & illumination correction: You can clearly see the 4 field of views based on the shading
   

2. The same 2x2 sites after the correction, now the shading is gone
   

Also, very nice flexibility of input handling and the overall ability to overwrite the Zarr files, which were the major goals for this issue! I'd say we have achieved both nice illumination correction, as well as building the necessary infrastructure here! 👏🏻🚀

the correct mapping between channels and files?
Had a quick chat with Gustavo. The matching should be correct, except for the A01_C04. Correct mapping should be (already pushed to main):

{
    "root_path_corr": "/data/active/fractal/Liberali/FractalTesting20220124/IlluminationCorrectionMatrices-Yokogawa/",
    "A04_C01": "220120_60xW_BP445_CH01.tif",
    "A03_C02": "220120_60xW_BP525_CH02.tif",
    "A02_C03": "220120_60xW_BP600_CH03.tif",
    "A01_C04": "220120_60xW_BP676_CH04.tif"
}

@tcompa Do you want to run the FMI dataset as a test now? If that looks good, I'd say with close this issue :)

---

• Side note: These were 16bit illumination profile and ran without an issue, thus showing that the flexibility we were aiming for regarding input format is working

[tcompa]

Great, thanks!

@tcompa Do you want to run the FMI dataset as a test now? If that looks good, I'd say with close this issue :)

Of course. I ran it this afternoon (before this commit), but I'll do it again and then close the issue.

[tcompa]

Here's a result for FMI 2-wells dataset:

Closing (for now).

[jluethi]

I tested the illumination correction on the 23 well dataset and looks very nice there as well :)
(using the UZH illumination profiles I created based on the TissueMaps illumination statistics)

Without illumination correction:

With illumination correction: