Add tool for analysing images using Bioimage AI models (#1391)

* add model inference tool * add tool tests and test data * fix linting issues * Update tools/bioimaging/bioimage_inference.xml Fix suggestion in name Co-authored-by: Leonid Kostrykin <[email protected]> * Update tools/bioimaging/bioimage_inference.xml Fix suggestion to include edam annotation Co-authored-by: Leonid Kostrykin <[email protected]> * update * add model size * fix dymanic input shapes * remove comments * upate * fix linting error * replace test files * update test files * add image size to test * fix remove comments * Apply suggestions from code review Fix review comments Co-authored-by: Beatriz Serrano-Solano <[email protected]> Co-authored-by: Leonid Kostrykin <[email protected]> * use correct model name * use original value of predicted matrix * add support for png * add creator * update bioimage name * fix review comments * Apply suggestions from code review Fix review comments Co-authored-by: Leonid Kostrykin <[email protected]> * fix review comments * Apply suggestions from code review Fix review comments Co-authored-by: Leonid Kostrykin <[email protected]> * Apply suggestions from code review Add missing articles and replace shape by size Co-authored-by: Leonid Kostrykin <[email protected]> * Apply suggestions from code review Fix review Co-authored-by: Björn Grüning <[email protected]> * fix review * restore tool version prefix * Apply suggestions from code review Co-authored-by: Leonid Kostrykin <[email protected]> --------- Co-authored-by: Leonid Kostrykin <[email protected]> Co-authored-by: Beatriz Serrano-Solano <[email protected]> Co-authored-by: Björn Grüning <[email protected]>
bgruening · Aug 2, 2024 · 57f4673 · 57f4673
1 parent 45375b6
commit 57f4673
Show file tree

Hide file tree

Showing 6 changed files with 177 additions and 0 deletions.
diff --git a/tools/bioimaging/.shed.yml b/tools/bioimaging/.shed.yml
@@ -0,0 +1,11 @@
+name: bioimage_inference
+owner: bgruening
+description: "Load model from BioImage.IO and make inferences"
+homepage_url: https://github.com/bgruening/galaxytools
+long_description: 
+    Load model from BioImage.IO and make inferences
+remote_repository_url: https://github.com/bgruening/galaxytools/tree/recommendation_training/tools/bioimaging
+type: unrestricted
+categories:
+- Imaging
+
diff --git a/tools/bioimaging/bioimage_inference.xml b/tools/bioimaging/bioimage_inference.xml
@@ -0,0 +1,80 @@
+<tool id="bioimage_inference" name="Process image using a BioImage.IO model" version="@TOOL_VERSION@+galaxy@VERSION_SUFFIX@" profile="23.0">
+    <description>with PyTorch</description>
+    <macros>
+        <token name="@TOOL_VERSION@">2.3.1</token>
+        <token name="@VERSION_SUFFIX@">0</token>
+    </macros>
+    <creator>
+        <organization name="European Galaxy Team" url="https://galaxyproject.org/eu/" />
+        <person givenName="Anup" familyName="Kumar" email="[email protected]" />
+        <person givenName="Beatriz" familyName="Serrano-Solano" email="[email protected]" />
+        <person givenName="Leonid" familyName="Kostrykin" email="[email protected]" />
+    </creator>
+    <edam_operations>
+        <edam_operation>operation_3443</edam_operation>
+    </edam_operations>
+    <xrefs>
+        <xref type="bio.tools">pytorch</xref>
+        <xref type="biii">pytorch</xref>
+    </xrefs>
+    <requirements>
+        <requirement type="package" version="3.9.12">python</requirement>
+        <requirement type="package" version="@TOOL_VERSION@">pytorch</requirement>
+        <requirement type="package" version="0.18.1">torchvision</requirement>
+        <requirement type="package" version="2.34.2">imageio</requirement>
+    </requirements>
+    <version_command>echo "@VERSION@"</version_command>
+    <command detect_errors="aggressive">
+    <![CDATA[
+        python '$__tool_directory__/main.py'
+            --imaging_model '$input_imaging_model'
+            --image_file '$input_image_file'
+            --image_size '$input_image_input_size'
+    ]]>
+    </command>
+    <inputs>
+        <param name="input_imaging_model" type="data" format="zip" label="BioImage.IO model" help="Please upload a BioImage.IO model."/>
+        <param name="input_image_file" type="data" format="tiff,png" label="Input image" help="Please provide an input image for the analysis."/>
+        <param name="input_image_input_size" type="text" label="Size of the input image" help="Provide the size of the input image. See the chosen model's RDF file to find the correct input size. For example: for the BioImage.IO model MitochondriaEMSegmentationBoundaryModel, the input size is 256 x 256 x 32 x 1. Enter the size as 256,256,32,1."/>
+    </inputs>
+    <outputs>
+        <data format="tif" name="output_predicted_image" from_work_dir="output_predicted_image.tif" label="Predicted image"></data>
+        <data format="npy" name="output_predicted_image_matrix" from_work_dir="output_predicted_image_matrix.npy" label="Predicted image tensor"></data>
+    </outputs>
+    <tests>
+        <test>
+            <param name="input_imaging_model" value="input_imaging_model.zip" location="https://zenodo.org/api/records/6647674/files/weights-torchscript.pt/content"/>
+            <param name="input_image_file" value="input_image_file.tif" location="https://zenodo.org/api/records/6647674/files/sample_input_0.tif/content"/>
+            <param name="input_image_input_size" value="256,256,1,1"/>
+            <output name="output_predicted_image" file="output_nucleisegboundarymodel.tif" compare="sim_size" delta="100" />
+            <output name="output_predicted_image_matrix" file="output_nucleisegboundarymodel_matrix.npy" compare="sim_size" delta="100" />
+        </test>
+        <test>
+            <param name="input_imaging_model" value="input_imaging_model.zip" location="https://zenodo.org/api/records/6647674/files/weights-torchscript.pt/content"/>
+            <param name="input_image_file" value="input_nucleisegboundarymodel.png"/>
+            <param name="input_image_input_size" value="256,256,1,1"/>
+            <output name="output_predicted_image" file="output_nucleisegboundarymodel.tif" compare="sim_size" delta="100" />
+            <output name="output_predicted_image_matrix" file="output_nucleisegboundarymodel_matrix.npy" compare="sim_size" delta="100" />
+        </test>
+    </tests>
+    <help>
+        <![CDATA[
+            **What it does**
+
+            The tool takes a BioImage.IO model and an image (as TIF or PNG) to be analyzed. The analysis is performed by the model. The model is used to obtain a prediction of the result of the analysis, and the predicted image becomes available as a TIF file in the Galaxy history.
+
+            **Input files**
+            - BioImage.IO model: Add one of the model from Galaxy file uploader by choosing a "remote" file at "ML Models/bioimaging-models"
+            - Image to be analyzed: Provide an image as TIF/PNG file
+            - Provide the necessary input size for the model. This information can be found in the RDF file of each model (RDF file > config > test_information > inputs > size)
+
+            **Output files**
+            - Predicted image: Predicted image using the BioImage.IO model
+            - Predicted image matrix: Predicted image matrix in original dimensions
+        ]]>
+    </help>
+    <citations>
+        <citation type="doi">10.1145/3620665.3640366</citation>
+        <citation type="doi">10.1101/2022.06.07.495102</citation>
+    </citations>
+</tool>
diff --git a/tools/bioimaging/main.py b/tools/bioimaging/main.py
@@ -0,0 +1,86 @@
+"""
+Predict images using AI models from BioImage.IO
+"""
+
+import argparse
+
+import imageio
+import numpy as np
+import torch
+
+
+def find_dim_order(user_in_shape, input_image):
+    """
+    Find the correct order of input image's
+    shape. For a few models, the order of input size
+    mentioned in the RDF.yaml file is reversed compared
+    to the input image's original size. If it is reversed,
+    transpose the image to find correct order of image's
+    dimensions.
+    """
+    image_shape = list(input_image.shape)
+    # reverse the input shape provided from RDF.yaml file
+    correct_order = user_in_shape.split(",")[::-1]
+    # remove 1s from the original dimensions
+    correct_order = [int(i) for i in correct_order if i != "1"]
+    if (correct_order[0] == image_shape[-1]) and (correct_order != image_shape):
+        input_image = torch.tensor(input_image.transpose())
+    return input_image, correct_order
+
+
+if __name__ == "__main__":
+    arg_parser = argparse.ArgumentParser()
+    arg_parser.add_argument("-im", "--imaging_model", required=True, help="Input BioImage model")
+    arg_parser.add_argument("-ii", "--image_file", required=True, help="Input image file")
+    arg_parser.add_argument("-is", "--image_size", required=True, help="Input image file's size")
+
+    # get argument values
+    args = vars(arg_parser.parse_args())
+    model_path = args["imaging_model"]
+    input_image_path = args["image_file"]
+
+    # load all embedded images in TIF file
+    test_data = imageio.v3.imread(input_image_path, index="...")
+    test_data = np.squeeze(test_data)
+    test_data = test_data.astype(np.float32)
+
+    # assess the correct dimensions of TIF input image
+    input_image_shape = args["image_size"]
+    im_test_data, shape_vals = find_dim_order(input_image_shape, test_data)
+
+    # load model
+    model = torch.load(model_path)
+    model.eval()
+
+    # find the number of dimensions required by the model
+    target_dimension = 0
+    for param in model.named_parameters():
+        target_dimension = len(param[1].shape)
+        break
+    current_dimension = len(list(im_test_data.shape))
+
+    # update the dimensions of input image if the required image by
+    # the model is smaller
+    slices = tuple(slice(0, s_val) for s_val in shape_vals)
+
+    # apply the slices to the reshaped_input
+    im_test_data = im_test_data[slices]
+    exp_test_data = torch.tensor(im_test_data)
+
+    # expand input image's dimensions
+    for i in range(target_dimension - current_dimension):
+        exp_test_data = torch.unsqueeze(exp_test_data, i)
+
+    # make prediction
+    pred_data = model(exp_test_data)
+    pred_data_output = pred_data.detach().numpy()
+
+    # save original image matrix
+    np.save("output_predicted_image_matrix.npy", pred_data_output)
+
+    # post process predicted file to correctly save as TIF file
+    pred_data = torch.squeeze(pred_data)
+    pred_numpy = pred_data.detach().numpy()
+
+    # write predicted TIF image to file
+    imageio.v3.imwrite("output_predicted_image.tif", pred_numpy, extension=".tif")
diff --git a/tools/bioimaging/test-data/input_nucleisegboundarymodel.png b/tools/bioimaging/test-data/input_nucleisegboundarymodel.png
diff --git a/tools/bioimaging/test-data/output_nucleisegboundarymodel.tif b/tools/bioimaging/test-data/output_nucleisegboundarymodel.tif
diff --git a/tools/bioimaging/test-data/output_nucleisegboundarymodel_matrix.npy b/tools/bioimaging/test-data/output_nucleisegboundarymodel_matrix.npy