From cf0e69bf5d1a9d69b71c692b1b1a31b42ba8ae7f Mon Sep 17 00:00:00 2001
From: Qin Yu <qin.yu.95@outlook.com>
Date: Mon, 16 Dec 2024 22:37:10 +0100
Subject: [PATCH] feat: support arbitrary network output suck as Cellpose

with broken types and returns
---
 plantseg/functionals/prediction/prediction.py |   12 +-
 plantseg/tasks/prediction_tasks.py            |    2 +-
 test_biio.py                                  |   17 -
 test_zoo.ipynb                                | 2554 -----------------
 4 files changed, 8 insertions(+), 2577 deletions(-)
 delete mode 100644 test_biio.py
 delete mode 100644 test_zoo.ipynb

diff --git a/plantseg/functionals/prediction/prediction.py b/plantseg/functionals/prediction/prediction.py
index f6296ac0..5b6e3377 100644
--- a/plantseg/functionals/prediction/prediction.py
+++ b/plantseg/functionals/prediction/prediction.py
@@ -93,11 +93,13 @@ def biio_prediction(
     assert isinstance(sample_out, Sample)
     if len(sample_out.members) != 1:
         logger.warning("Model has more than one output tensor. PlantSeg does not support this yet.")
-    key = list(sample_out.members.keys())[0]
-    pmaps = sample_out.members[key].data.to_numpy()[0]
-    assert pmaps.ndim == 4, f"Expected 4D CZXY prediction from `biio_prediction()`, got {pmaps.ndim}D"
-
-    return pmaps
+    t = {i: o.transpose(['batch', 'channel', 'z', 'y', 'x']) for i, o in sample_out.members.items()}
+    pmaps = []
+    for i, bczyx in t.items():
+        for czyx in bczyx:
+            for zyx in czyx:
+                pmaps.append(zyx.data.to_numpy())
+    return pmaps  # FIXME: Wrong return type
 
 
 def unet_prediction(
diff --git a/plantseg/tasks/prediction_tasks.py b/plantseg/tasks/prediction_tasks.py
index 4f185615..7e9f6402 100644
--- a/plantseg/tasks/prediction_tasks.py
+++ b/plantseg/tasks/prediction_tasks.py
@@ -51,7 +51,7 @@ def unet_prediction_task(
         config_path=config_path,
         model_weights_path=model_weights_path,
     )
-    assert pmaps.ndim == 4, f"Expected 4D CZXY prediction, got {pmaps.ndim}D"
+    # assert pmaps.ndim == 4, f"Expected 4D CZXY prediction, got {pmaps.ndim}D"
 
     new_images = []
 
diff --git a/test_biio.py b/test_biio.py
deleted file mode 100644
index 1f8d2522..00000000
--- a/test_biio.py
+++ /dev/null
@@ -1,17 +0,0 @@
-import numpy as np
-from bioimageio.core.prediction import predict
-from bioimageio.core.sample import Sample
-from bioimageio.core.tensor import Tensor
-from bioimageio.spec.model.v0_5 import TensorId
-
-array = np.random.randint(0, 255, (2, 128, 128, 128), dtype=np.uint8)
-dims = ('c', 'z', 'y', 'x')
-sample = Sample(members={TensorId('a'): Tensor(array=array, dims=dims)}, stat={}, id='try')
-
-temp = predict(
-    # model='https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/rdf.yaml',
-    model='https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/emotional-cricket/1.1/files/rdf.yaml',
-    # model='/Users/qin/Downloads/efficient-chipmunk.yaml',
-    inputs=sample,
-    sample_id='sample',
-)
diff --git a/test_zoo.ipynb b/test_zoo.ipynb
deleted file mode 100644
index 5cf42488..00000000
--- a/test_zoo.ipynb
+++ /dev/null
@@ -1,2554 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO: P [MainThread] 2024-12-05 16:54:48,153 plantseg - Logger configured at initialisation. PlantSeg logger name: plantseg\n"
-     ]
-    }
-   ],
-   "source": [
-    "from plantseg.core.zoo import model_zoo"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "PlantSeg version: 2.0.0a7\n",
-      "PyTorch version: 2.2.2\n"
-     ]
-    }
-   ],
-   "source": [
-    "import torch\n",
-    "\n",
-    "from plantseg import __version__\n",
-    "\n",
-    "print(f\"PlantSeg version: {__version__.__version__}\")\n",
-    "print(f\"PyTorch version: {torch.__version__}\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model_zoo.refresh_bioimageio_zoo_urls()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "<style scoped>\n",
-       "    .dataframe tbody tr th:only-of-type {\n",
-       "        vertical-align: middle;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe tbody tr th {\n",
-       "        vertical-align: top;\n",
-       "    }\n",
-       "\n",
-       "    .dataframe thead th {\n",
-       "        text-align: right;\n",
-       "    }\n",
-       "</style>\n",
-       "<table border=\"1\" class=\"dataframe\">\n",
-       "  <thead>\n",
-       "    <tr style=\"text-align: right;\">\n",
-       "      <th></th>\n",
-       "      <th>name</th>\n",
-       "      <th>url</th>\n",
-       "      <th>path</th>\n",
-       "      <th>description</th>\n",
-       "      <th>resolution</th>\n",
-       "      <th>dimensionality</th>\n",
-       "      <th>modality</th>\n",
-       "      <th>recommended_patch_size</th>\n",
-       "      <th>output_type</th>\n",
-       "      <th>doi</th>\n",
-       "      <th>added_by</th>\n",
-       "      <th>name_display</th>\n",
-       "      <th>rdf_source</th>\n",
-       "      <th>supported</th>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>id</th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "      <th></th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>affable-shark</th>\n",
-       "      <td>NucleiSegmentationBoundaryModel</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>affable-shark            : NucleiSegmentationB...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>chatty-frog</th>\n",
-       "      <td>StarDist H&amp;E Nuclei Segmentation</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>chatty-frog              : StarDist H&amp;E Nuclei...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>hiding-tiger</th>\n",
-       "      <td>LiveCellSegmentationBoundaryModel</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>hiding-tiger             : LiveCellSegmentatio...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>impartial-shrimp</th>\n",
-       "      <td>Neuron Segmentation in EM (Membrane Prediction)</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>impartial-shrimp         : Neuron Segmentation...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>discreet-rooster</th>\n",
-       "      <td>Pancreatic Phase Contrast Cell Segmentation (U...</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>discreet-rooster         : Pancreatic Phase Co...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>stupendous-sheep</th>\n",
-       "      <td>(Empanada) 2D Instance Mitochondrial Segmentat...</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>stupendous-sheep         : (Empanada) 2D Insta...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>pioneering-goat</th>\n",
-       "      <td>UniFMIRProjectionOnFlyWing</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>pioneering-goat          : UniFMIRProjectionOn...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>greedy-shark</th>\n",
-       "      <td>UniFMIRVolumetricReconstructionOnVCD</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>greedy-shark             : UniFMIRVolumetricRe...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>lucky-fox</th>\n",
-       "      <td>UniFMIRIsotropicReconstructionOnLiver</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>lucky-fox                : UniFMIRIsotropicRec...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>frank-water-buffalo</th>\n",
-       "      <td>UniFMIRDenoiseOnPlanaria</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>None</td>\n",
-       "      <td>bioimage.io</td>\n",
-       "      <td>frank-water-buffalo      : UniFMIRDenoiseOnPla...</td>\n",
-       "      <td>https://uk1s3.embassy.ebi.ac.uk/public-dataset...</td>\n",
-       "      <td>False</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>68 rows × 14 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "                                                                  name   url  \\\n",
-       "id                                                                             \n",
-       "affable-shark                          NucleiSegmentationBoundaryModel  None   \n",
-       "chatty-frog                           StarDist H&E Nuclei Segmentation  None   \n",
-       "hiding-tiger                         LiveCellSegmentationBoundaryModel  None   \n",
-       "impartial-shrimp       Neuron Segmentation in EM (Membrane Prediction)  None   \n",
-       "discreet-rooster     Pancreatic Phase Contrast Cell Segmentation (U...  None   \n",
-       "...                                                                ...   ...   \n",
-       "stupendous-sheep     (Empanada) 2D Instance Mitochondrial Segmentat...  None   \n",
-       "pioneering-goat                             UniFMIRProjectionOnFlyWing  None   \n",
-       "greedy-shark                      UniFMIRVolumetricReconstructionOnVCD  None   \n",
-       "lucky-fox                        UniFMIRIsotropicReconstructionOnLiver  None   \n",
-       "frank-water-buffalo                           UniFMIRDenoiseOnPlanaria  None   \n",
-       "\n",
-       "                     path description resolution dimensionality modality  \\\n",
-       "id                                                                         \n",
-       "affable-shark        None        None       None           None     None   \n",
-       "chatty-frog          None        None       None           None     None   \n",
-       "hiding-tiger         None        None       None           None     None   \n",
-       "impartial-shrimp     None        None       None           None     None   \n",
-       "discreet-rooster     None        None       None           None     None   \n",
-       "...                   ...         ...        ...            ...      ...   \n",
-       "stupendous-sheep     None        None       None           None     None   \n",
-       "pioneering-goat      None        None       None           None     None   \n",
-       "greedy-shark         None        None       None           None     None   \n",
-       "lucky-fox            None        None       None           None     None   \n",
-       "frank-water-buffalo  None        None       None           None     None   \n",
-       "\n",
-       "                    recommended_patch_size output_type   doi     added_by  \\\n",
-       "id                                                                          \n",
-       "affable-shark                         None        None  None  bioimage.io   \n",
-       "chatty-frog                           None        None  None  bioimage.io   \n",
-       "hiding-tiger                          None        None  None  bioimage.io   \n",
-       "impartial-shrimp                      None        None  None  bioimage.io   \n",
-       "discreet-rooster                      None        None  None  bioimage.io   \n",
-       "...                                    ...         ...   ...          ...   \n",
-       "stupendous-sheep                      None        None  None  bioimage.io   \n",
-       "pioneering-goat                       None        None  None  bioimage.io   \n",
-       "greedy-shark                          None        None  None  bioimage.io   \n",
-       "lucky-fox                             None        None  None  bioimage.io   \n",
-       "frank-water-buffalo                   None        None  None  bioimage.io   \n",
-       "\n",
-       "                                                          name_display  \\\n",
-       "id                                                                       \n",
-       "affable-shark        affable-shark            : NucleiSegmentationB...   \n",
-       "chatty-frog          chatty-frog              : StarDist H&E Nuclei...   \n",
-       "hiding-tiger         hiding-tiger             : LiveCellSegmentatio...   \n",
-       "impartial-shrimp     impartial-shrimp         : Neuron Segmentation...   \n",
-       "discreet-rooster     discreet-rooster         : Pancreatic Phase Co...   \n",
-       "...                                                                ...   \n",
-       "stupendous-sheep     stupendous-sheep         : (Empanada) 2D Insta...   \n",
-       "pioneering-goat      pioneering-goat          : UniFMIRProjectionOn...   \n",
-       "greedy-shark         greedy-shark             : UniFMIRVolumetricRe...   \n",
-       "lucky-fox            lucky-fox                : UniFMIRIsotropicRec...   \n",
-       "frank-water-buffalo  frank-water-buffalo      : UniFMIRDenoiseOnPla...   \n",
-       "\n",
-       "                                                            rdf_source  \\\n",
-       "id                                                                       \n",
-       "affable-shark        https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "chatty-frog          https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "hiding-tiger         https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "impartial-shrimp     https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "discreet-rooster     https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "...                                                                ...   \n",
-       "stupendous-sheep     https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "pioneering-goat      https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "greedy-shark         https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "lucky-fox            https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "frank-water-buffalo  https://uk1s3.embassy.ebi.ac.uk/public-dataset...   \n",
-       "\n",
-       "                     supported  \n",
-       "id                              \n",
-       "affable-shark            False  \n",
-       "chatty-frog              False  \n",
-       "hiding-tiger             False  \n",
-       "impartial-shrimp         False  \n",
-       "discreet-rooster         False  \n",
-       "...                        ...  \n",
-       "stupendous-sheep         False  \n",
-       "pioneering-goat          False  \n",
-       "greedy-shark             False  \n",
-       "lucky-fox                False  \n",
-       "frank-water-buffalo      False  \n",
-       "\n",
-       "[68 rows x 14 columns]"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_zoo.models_bioimageio"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['affable-shark',\n",
-       " 'affectionate-cow',\n",
-       " 'ambitious-ant',\n",
-       " 'ambitious-sloth',\n",
-       " 'amiable-crocodile',\n",
-       " 'charismatic-whale',\n",
-       " 'chatty-frog',\n",
-       " 'committed-turkey',\n",
-       " 'conscientious-seashell',\n",
-       " 'courteous-otter',\n",
-       " 'creative-panda',\n",
-       " 'dazzling-spider',\n",
-       " 'decisive-panda',\n",
-       " 'determined-chipmunk',\n",
-       " 'determined-hedgehog',\n",
-       " 'diplomatic-bug',\n",
-       " 'discreet-rooster',\n",
-       " 'dynamic-t-rex',\n",
-       " 'easy-going-sauropod',\n",
-       " 'efficient-chipmunk',\n",
-       " 'emotional-cricket',\n",
-       " 'faithful-chicken',\n",
-       " 'famous-fish',\n",
-       " 'fearless-crab',\n",
-       " 'frank-water-buffalo',\n",
-       " 'greedy-shark',\n",
-       " 'greedy-whale',\n",
-       " 'happy-elephant',\n",
-       " 'hiding-blowfish',\n",
-       " 'hiding-tiger',\n",
-       " 'humorous-crab',\n",
-       " 'humorous-fox',\n",
-       " 'humorous-owl',\n",
-       " 'idealistic-rat',\n",
-       " 'impartial-shark',\n",
-       " 'impartial-shrimp',\n",
-       " 'independent-shrimp',\n",
-       " 'joyful-deer',\n",
-       " 'kind-seashell',\n",
-       " 'laid-back-lobster',\n",
-       " 'loyal-parrot',\n",
-       " 'loyal-squid',\n",
-       " 'lucky-fox',\n",
-       " 'modest-octopus',\n",
-       " 'naked-microbe',\n",
-       " 'nice-peacock',\n",
-       " 'noisy-fish',\n",
-       " 'noisy-hedgehog',\n",
-       " 'noisy-ox',\n",
-       " 'non-judgemental-eagle',\n",
-       " 'organized-badger',\n",
-       " 'organized-cricket',\n",
-       " 'passionate-t-rex',\n",
-       " 'philosophical-panda',\n",
-       " 'pioneering-goat',\n",
-       " 'pioneering-rhino',\n",
-       " 'placid-llama',\n",
-       " 'polite-pig',\n",
-       " 'powerful-chipmunk',\n",
-       " 'powerful-fish',\n",
-       " 'resourceful-lizard',\n",
-       " 'shivering-raccoon',\n",
-       " 'straightforward-crocodile',\n",
-       " 'stupendous-sheep',\n",
-       " 'thoughtful-turtle',\n",
-       " 'wild-rhino',\n",
-       " 'wild-whale',\n",
-       " 'willing-hedgehog']"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# get all index\n",
-    "\n",
-    "sorted(model_zoo.models_bioimageio.index.to_list())"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['efficient-chipmunk',\n",
-       " 'emotional-cricket',\n",
-       " 'loyal-squid',\n",
-       " 'noisy-fish',\n",
-       " 'passionate-t-rex',\n",
-       " 'pioneering-rhino',\n",
-       " 'powerful-fish',\n",
-       " 'thoughtful-turtle']"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "# get all index where `supported` is True:\n",
-    "\n",
-    "sorted(model_zoo.models_bioimageio[model_zoo.models_bioimageio[\"supported\"]].index.to_list())"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "'efficient-chipmunk' in model_zoo.models_bioimageio.index"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "name                                          PlantSeg Plant Nuclei 3D UNet\n",
-       "url                                                                    None\n",
-       "path                                                                   None\n",
-       "description                                                            None\n",
-       "resolution                                                             None\n",
-       "dimensionality                                                         None\n",
-       "modality                                                               None\n",
-       "recommended_patch_size                                                 None\n",
-       "output_type                                                            None\n",
-       "doi                                                                    None\n",
-       "added_by                                                        bioimage.io\n",
-       "name_display              efficient-chipmunk       : PlantSeg Plant Nucl...\n",
-       "rdf_source                https://uk1s3.embassy.ebi.ac.uk/public-dataset...\n",
-       "supported                                                              True\n",
-       "Name: efficient-chipmunk, dtype: object"
-      ]
-     },
-     "execution_count": 8,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_zoo.models_bioimageio.loc['efficient-chipmunk']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Url('https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/efficient-chipmunk/1/files/rdf.yaml')"
-      ]
-     },
-     "execution_count": 9,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_zoo.models_bioimageio.loc['efficient-chipmunk']['rdf_source']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Url('https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/emotional-cricket/1.1/files/rdf.yaml')"
-      ]
-     },
-     "execution_count": 20,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_zoo.models_bioimageio.loc['emotional-cricket']['rdf_source']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "id\n",
-       "affable-shark                            NucleiSegmentationBoundaryModel\n",
-       "chatty-frog                             StarDist H&E Nuclei Segmentation\n",
-       "hiding-tiger                           LiveCellSegmentationBoundaryModel\n",
-       "impartial-shrimp         Neuron Segmentation in EM (Membrane Prediction)\n",
-       "discreet-rooster       Pancreatic Phase Contrast Cell Segmentation (U...\n",
-       "                                             ...                        \n",
-       "stupendous-sheep       (Empanada) 2D Instance Mitochondrial Segmentat...\n",
-       "pioneering-goat                               UniFMIRProjectionOnFlyWing\n",
-       "greedy-shark                        UniFMIRVolumetricReconstructionOnVCD\n",
-       "lucky-fox                          UniFMIRIsotropicReconstructionOnLiver\n",
-       "frank-water-buffalo                             UniFMIRDenoiseOnPlanaria\n",
-       "Name: name, Length: 68, dtype: object"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_zoo.models_bioimageio['name']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "pydantic_core._pydantic_core.Url"
-      ]
-     },
-     "execution_count": 11,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "type(model_zoo.models_bioimageio.at['efficient-chipmunk', 'rdf_source'])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "pydantic_core._pydantic_core.Url"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "type(model_zoo.models_bioimageio.loc['efficient-chipmunk']['rdf_source'])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import numpy as np\n",
-    "from bioimageio.core.prediction import predict\n",
-    "from bioimageio.core.sample import Sample\n",
-    "from bioimageio.core.tensor import Tensor\n",
-    "from bioimageio.spec.model.v0_5 import TensorId"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "array = np.random.randint(0, 255, (128, 128, 128), dtype=np.uint8)\n",
-    "dims = ('z', 'y', 'x')\n",
-    "sample = Sample(members={TensorId('a'): Tensor(array, dims)}, stat={}, id='try')\n",
-    "# sample.members[TensorId('a')].data"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Url('https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/rdf.yaml')"
-      ]
-     },
-     "execution_count": 19,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model_zoo.models_bioimageio.at['philosophical-panda', 'rdf_source']"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/rdf.yaml' to file '/Users/qin/Library/Caches/bioimageio/520a69782c7dafb6478e43ebcc4679b0-rdf.yaml'.\n",
-      "100%|█████████████████████████████████████| 13.5k/13.5k [00:00<00:00, 35.7MB/s]\n",
-      "SHA256 hash of downloaded file: bbad75237ecf4f9d9f6259b13b97fc01b5cbeb4e3ea672e72826608d68197a32\n",
-      "Use this value as the 'known_hash' argument of 'pooch.retrieve' to ensure that the file hasn't changed if it is downloaded again in the future.\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/README.md' to file '/Users/qin/Library/Caches/bioimageio/42620dae3e3850cefbf0475c7cf590dd-README.md'.\n",
-      "100%|██████████████████████████████████████████| 431/431 [00:00<00:00, 658kB/s]\n",
-      "SHA256 hash of downloaded file: fc6e1292ca309bedaca504260cecc9a7bc9f26e9328eb7a051f82a2ceec475e3\n",
-      "Use this value as the 'known_hash' argument of 'pooch.retrieve' to ensure that the file hasn't changed if it is downloaded again in the future.\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_input.npy' to file '/Users/qin/Library/Caches/bioimageio/101853864f8c8e986b2819c9ac44d0f9-test_input.npy'.\n",
-      "100%|█████████████████████████████████████| 5.53M/5.53M [00:00<00:00, 8.91GB/s]\n",
-      "computing SHA256 of 101853864f8c8e986b2819c9ac44d0f9-test_input.npy (result: 6810255f5b5260fe39153f2192bedf30d9899ec4e770976b7813116c467579f0): 100%|██████████| 5529728/5529728 [00:00<00:00, 1220895945.11it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_output.npy' to file '/Users/qin/Library/Caches/bioimageio/57b964a123db3ff8400bcce3ef902b18-test_output.npy'.\n",
-      "100%|█████████████████████████████████████| 8.29M/8.29M [00:00<00:00, 11.9GB/s]\n",
-      "computing SHA256 of 57b964a123db3ff8400bcce3ef902b18-test_output.npy (result: d802e3024da80bff93a9ec50fbe50b9c3946534aab1b60b911511111a8e2dbca): 100%|██████████| 8294528/8294528 [00:00<00:00, 1690711569.64it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_style.npy' to file '/Users/qin/Library/Caches/bioimageio/0b712fa8abf8707021a71747e726bf7c-test_style.npy'.\n",
-      "100%|█████████████████████████████████████| 76.9k/76.9k [00:00<00:00, 71.6MB/s]\n",
-      "computing SHA256 of 0b712fa8abf8707021a71747e726bf7c-test_style.npy (result: ab464b406f9050561b40f7d76700ab5edf3aca97e31fe9a6069a51aeeca8bc81): 100%|██████████| 76928/76928 [00:00<00:00, 106734838.94it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_downsampled_0.npy' to file '/Users/qin/Library/Caches/bioimageio/d0abd68ef3844b0d6fbed811e0fed878-test_downsampled_0.npy'.\n",
-      "100%|█████████████████████████████████████| 88.5M/88.5M [00:00<00:00, 82.3GB/s]\n",
-      "computing SHA256 of d0abd68ef3844b0d6fbed811e0fed878-test_downsampled_0.npy (result: 67df53fb440e94dbb9c8e4003dcbde158646a7975c4878cacdd251e1fcfb4225): 100%|██████████| 88473728/88473728 [00:00<00:00, 2235146402.88it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_downsampled_1.npy' to file '/Users/qin/Library/Caches/bioimageio/3765cac1d92a49daf0d6ec949919aeb1-test_downsampled_1.npy'.\n",
-      "100%|█████████████████████████████████████| 44.2M/44.2M [00:00<00:00, 52.2GB/s]\n",
-      "computing SHA256 of 3765cac1d92a49daf0d6ec949919aeb1-test_downsampled_1.npy (result: cb4addbd763d96731ebd18ed001b87ab7195ec9198f01a753a363a06c27bfb1c): 100%|██████████| 44236928/44236928 [00:00<00:00, 2150700977.83it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_downsampled_2.npy' to file '/Users/qin/Library/Caches/bioimageio/7251078a2afa8713384a3103878dd09d-test_downsampled_2.npy'.\n",
-      "100%|█████████████████████████████████████| 22.1M/22.1M [00:00<00:00, 24.9GB/s]\n",
-      "computing SHA256 of 7251078a2afa8713384a3103878dd09d-test_downsampled_2.npy (result: 9c0225b94d84fcc3adfb9a73eef1303d6adb318b57a5a801e0e2e1638b458e72): 100%|██████████| 22118528/22118528 [00:00<00:00, 2035139420.08it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/test_downsampled_3.npy' to file '/Users/qin/Library/Caches/bioimageio/e434cdc3ea3e7ecfb752cdc001617875-test_downsampled_3.npy'.\n",
-      "100%|█████████████████████████████████████| 11.1M/11.1M [00:00<00:00, 15.7GB/s]\n",
-      "computing SHA256 of e434cdc3ea3e7ecfb752cdc001617875-test_downsampled_3.npy (result: 1ea789ff37d47197c847b585799f7d063e7592b0c5e9c3094fd0e3ac209b7fc2): 100%|██████████| 11059328/11059328 [00:00<00:00, 1367718816.68it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/cpnet_wrapper.py' to file '/Users/qin/Library/Caches/bioimageio/00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py'.\n",
-      "100%|█████████████████████████████████████| 11.1k/11.1k [00:00<00:00, 18.9MB/s]\n",
-      "computing SHA256 of 00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py (result: b8b947cdd0ea8f5b98bd7be5f12f38bb1ea1ebe0b455c62d9a6389cd21d134bf): 100%|██████████| 11053/11053 [00:00<00:00, 15037185.25it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/cp_state_dict_1135_gold.pth' to file '/Users/qin/Library/Caches/bioimageio/8dbb20d5a3cb3a3dfdb5101a671861ce-cp_state_dict_1135_gold.pth'.\n",
-      "100%|█████████████████████████████████████| 26.6M/26.6M [00:00<00:00, 43.3GB/s]\n",
-      "computing SHA256 of 8dbb20d5a3cb3a3dfdb5101a671861ce-cp_state_dict_1135_gold.pth (result: 26c277f3b8f6ca5aab30b4b0a832601aea60183cbed1c2333576f4135a643eb2): 100%|██████████| 26556687/26556687 [00:00<00:00, 2103662363.99it/s]\n",
-      "Downloading data from 'https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/cp_traced_1135_gold.pt' to file '/Users/qin/Library/Caches/bioimageio/17fee110c39ccad7c3cb36d00d2fdd2c-cp_traced_1135_gold.pt'.\n",
-      "100%|█████████████████████████████████████| 26.8M/26.8M [00:00<00:00, 26.3GB/s]\n",
-      "computing SHA256 of 17fee110c39ccad7c3cb36d00d2fdd2c-cp_traced_1135_gold.pt (result: f61bae146ab522902350eadda1d509ac1037726fe6d7fb63f6a8a314021d63e7): 100%|██████████| 26812339/26812339 [00:00<00:00, 2032773000.69it/s]\n"
-     ]
-    },
-    {
-     "ename": "AttributeError",
-     "evalue": "'NoneType' object has no attribute 'dim'",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[15], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m temp \u001b[38;5;241m=\u001b[39m \u001b[43mpredict\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      2\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# model=model_zoo.models_bioimageio.at['emotional-cricket', 'rdf_source'],\u001b[39;49;00m\n\u001b[1;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# model=model_zoo.models_bioimageio.at['efficient-chipmunk', 'rdf_source'],\u001b[39;49;00m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmodel\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmodel_zoo\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodels_bioimageio\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mat\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mphilosophical-panda\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43mrdf_source\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;66;43;03m# model='/Users/qin/Downloads/rdf.yaml',\u001b[39;49;00m\n\u001b[1;32m      6\u001b[0m \u001b[43m    \u001b[49m\u001b[43minputs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msample\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m    \u001b[49m\u001b[43msample_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[38;5;124;43msample\u001b[39;49m\u001b[38;5;124;43m'\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/bioimageio/core/prediction.py:114\u001b[0m, in \u001b[0;36mpredict\u001b[0;34m(model, inputs, sample_id, blocksize_parameter, input_block_shape, skip_preprocessing, skip_postprocessing, save_output_path)\u001b[0m\n\u001b[1;32m    107\u001b[0m     output \u001b[38;5;241m=\u001b[39m pp\u001b[38;5;241m.\u001b[39mpredict_sample_with_blocking(\n\u001b[1;32m    108\u001b[0m         sample,\n\u001b[1;32m    109\u001b[0m         skip_preprocessing\u001b[38;5;241m=\u001b[39mskip_preprocessing,\n\u001b[1;32m    110\u001b[0m         skip_postprocessing\u001b[38;5;241m=\u001b[39mskip_postprocessing,\n\u001b[1;32m    111\u001b[0m         ns\u001b[38;5;241m=\u001b[39mblocksize_parameter,\n\u001b[1;32m    112\u001b[0m     )\n\u001b[1;32m    113\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 114\u001b[0m     output \u001b[38;5;241m=\u001b[39m \u001b[43mpp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpredict_sample_without_blocking\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    115\u001b[0m \u001b[43m        \u001b[49m\u001b[43msample\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    116\u001b[0m \u001b[43m        \u001b[49m\u001b[43mskip_preprocessing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mskip_preprocessing\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    117\u001b[0m \u001b[43m        \u001b[49m\u001b[43mskip_postprocessing\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mskip_postprocessing\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    118\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    119\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m save_output_path:\n\u001b[1;32m    120\u001b[0m     save_sample(save_output_path, output)\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/bioimageio/core/_prediction_pipeline.py:160\u001b[0m, in \u001b[0;36mPredictionPipeline.predict_sample_without_blocking\u001b[0;34m(self, sample, skip_preprocessing, skip_postprocessing)\u001b[0m\n\u001b[1;32m    152\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m skip_preprocessing:\n\u001b[1;32m    153\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mapply_preprocessing(sample)\n\u001b[1;32m    155\u001b[0m output \u001b[38;5;241m=\u001b[39m Sample(\n\u001b[1;32m    156\u001b[0m     members\u001b[38;5;241m=\u001b[39m{\n\u001b[1;32m    157\u001b[0m         out_id: out\n\u001b[1;32m    158\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m out_id, out \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(\n\u001b[1;32m    159\u001b[0m             \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_output_ids,\n\u001b[0;32m--> 160\u001b[0m             \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_adapter\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    161\u001b[0m \u001b[43m                \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43msample\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmembers\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mget\u001b[49m\u001b[43m(\u001b[49m\u001b[43min_id\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43min_id\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_input_ids\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    162\u001b[0m \u001b[43m            \u001b[49m\u001b[43m)\u001b[49m,\n\u001b[1;32m    163\u001b[0m         )\n\u001b[1;32m    164\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m out \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    165\u001b[0m     },\n\u001b[1;32m    166\u001b[0m     stat\u001b[38;5;241m=\u001b[39msample\u001b[38;5;241m.\u001b[39mstat,\n\u001b[1;32m    167\u001b[0m     \u001b[38;5;28mid\u001b[39m\u001b[38;5;241m=\u001b[39msample\u001b[38;5;241m.\u001b[39mid,\n\u001b[1;32m    168\u001b[0m )\n\u001b[1;32m    169\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m skip_postprocessing:\n\u001b[1;32m    170\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mapply_postprocessing(output)\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/bioimageio/core/model_adapters/_pytorch_model_adapter.py:71\u001b[0m, in \u001b[0;36mPytorchModelAdapter.forward\u001b[0;34m(self, *input_tensors)\u001b[0m\n\u001b[1;32m     60\u001b[0m tensors \u001b[38;5;241m=\u001b[39m [\n\u001b[1;32m     61\u001b[0m     (\n\u001b[1;32m     62\u001b[0m         \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     68\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m t \u001b[38;5;129;01min\u001b[39;00m tensors\n\u001b[1;32m     69\u001b[0m ]\n\u001b[1;32m     70\u001b[0m result: Union[Tuple[Any, \u001b[38;5;241m.\u001b[39m\u001b[38;5;241m.\u001b[39m\u001b[38;5;241m.\u001b[39m], List[Any], Any]\n\u001b[0;32m---> 71\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_network\u001b[49m\u001b[43m(\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# pyright: ignore[reportUnknownVariableType]\u001b[39;49;00m\n\u001b[1;32m     72\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtensors\u001b[49m\n\u001b[1;32m     73\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     74\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(result, (\u001b[38;5;28mtuple\u001b[39m, \u001b[38;5;28mlist\u001b[39m)):\n\u001b[1;32m     75\u001b[0m     result \u001b[38;5;241m=\u001b[39m [result]\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m~/Library/Caches/bioimageio/00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py:283\u001b[0m, in \u001b[0;36mCPnetBioImageIO.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m    273\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[1;32m    274\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    275\u001b[0m \u001b[38;5;124;03m    Perform a forward pass of the CPnet model and return unpacked tensors.\u001b[39;00m\n\u001b[1;32m    276\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    281\u001b[0m \u001b[38;5;124;03m        tuple: A tuple containing the output tensor, style tensor, and downsampled tensors.\u001b[39;00m\n\u001b[1;32m    282\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 283\u001b[0m     output_tensor, style_tensor, downsampled_tensors \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mforward\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    284\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m output_tensor, style_tensor, \u001b[38;5;241m*\u001b[39mdownsampled_tensors\n",
-      "File \u001b[0;32m~/Library/Caches/bioimageio/00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py:207\u001b[0m, in \u001b[0;36mCPnet.forward\u001b[0;34m(self, data)\u001b[0m\n\u001b[1;32m    205\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmkldnn:\n\u001b[1;32m    206\u001b[0m     data \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39mto_mkldnn()\n\u001b[0;32m--> 207\u001b[0m T0 \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdownsample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    208\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmkldnn:\n\u001b[1;32m    209\u001b[0m     style \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmake_style(T0[\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m]\u001b[38;5;241m.\u001b[39mto_dense())\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m~/Library/Caches/bioimageio/00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py:60\u001b[0m, in \u001b[0;36mdownsample.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     58\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     59\u001b[0m         y \u001b[38;5;241m=\u001b[39m x\n\u001b[0;32m---> 60\u001b[0m     xd\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdown\u001b[49m\u001b[43m[\u001b[49m\u001b[43mn\u001b[49m\u001b[43m]\u001b[49m\u001b[43m(\u001b[49m\u001b[43my\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     61\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m xd\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m~/Library/Caches/bioimageio/00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py:37\u001b[0m, in \u001b[0;36mresdown.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     36\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x):\n\u001b[0;32m---> 37\u001b[0m     x \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mproj\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv[\u001b[38;5;241m1\u001b[39m](\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv[\u001b[38;5;241m0\u001b[39m](x))\n\u001b[1;32m     38\u001b[0m     x \u001b[38;5;241m=\u001b[39m x \u001b[38;5;241m+\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv[\u001b[38;5;241m3\u001b[39m](\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconv[\u001b[38;5;241m2\u001b[39m](x))\n\u001b[1;32m     39\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m x\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/container.py:217\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    215\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[1;32m    216\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m:\n\u001b[0;32m--> 217\u001b[0m         \u001b[38;5;28minput\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[43mmodule\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    218\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28minput\u001b[39m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1511\u001b[0m, in \u001b[0;36mModule._wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1509\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_compiled_call_impl(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[1;32m   1510\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1511\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/module.py:1520\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1515\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1516\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks\n\u001b[1;32m   1518\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1519\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1520\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1522\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1523\u001b[0m     result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/batchnorm.py:142\u001b[0m, in \u001b[0;36m_BatchNorm.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    141\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m: Tensor) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Tensor:\n\u001b[0;32m--> 142\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_check_input_dim\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    144\u001b[0m     \u001b[38;5;66;03m# exponential_average_factor is set to self.momentum\u001b[39;00m\n\u001b[1;32m    145\u001b[0m     \u001b[38;5;66;03m# (when it is available) only so that it gets updated\u001b[39;00m\n\u001b[1;32m    146\u001b[0m     \u001b[38;5;66;03m# in ONNX graph when this node is exported to ONNX.\u001b[39;00m\n\u001b[1;32m    147\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmomentum \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
-      "File \u001b[0;32m~/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/torch/nn/modules/batchnorm.py:419\u001b[0m, in \u001b[0;36mBatchNorm2d._check_input_dim\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    418\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_check_input_dim\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;28minput\u001b[39m):\n\u001b[0;32m--> 419\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;43minput\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdim\u001b[49m() \u001b[38;5;241m!=\u001b[39m \u001b[38;5;241m4\u001b[39m:\n\u001b[1;32m    420\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mexpected 4D input (got \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28minput\u001b[39m\u001b[38;5;241m.\u001b[39mdim()\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124mD input)\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
-      "\u001b[0;31mAttributeError\u001b[0m: 'NoneType' object has no attribute 'dim'"
-     ]
-    }
-   ],
-   "source": [
-    "temp = predict(\n",
-    "    # model=model_zoo.models_bioimageio.at['emotional-cricket', 'rdf_source'],\n",
-    "    # model=model_zoo.models_bioimageio.at['efficient-chipmunk', 'rdf_source'],\n",
-    "    model=model_zoo.models_bioimageio.at['philosophical-panda', 'rdf_source'],\n",
-    "    # model='/Users/qin/Downloads/rdf.yaml',\n",
-    "    inputs=sample,\n",
-    "    sample_id='sample',\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "((1, 32, 64, 64), (64, 64, 32, 1))"
-      ]
-     },
-     "execution_count": 23,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "np.random.rand(1, 32, 64, 64).shape, np.random.rand(64, 64, 32, 1).shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from plantseg.functionals.dataprocessing.dataprocessing import ImageLayout"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "typing.Literal['ZYX', 'YX', 'CZYX', 'CYX']"
-      ]
-     },
-     "execution_count": 2,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "ImageLayout"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "('Z', 'Y', 'X')"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "tuple('ZYX')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO: P [MainThread] 2024-12-14 01:09:58,630 plantseg - Logger configured at initialisation. PlantSeg logger name: plantseg\n"
-     ]
-    }
-   ],
-   "source": [
-    "from pathlib import Path\n",
-    "from plantseg.functionals.prediction.prediction import biio_prediction\n",
-    "from plantseg.io import smart_load"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(453, 800, 800)"
-      ]
-     },
-     "execution_count": 5,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "raw = smart_load(\n",
-    "    Path('/Users/qin/Documents/Work/Side_rAP2_16LDs_SAM4_nuclei_Z0.400_X0.291_Y0.291_Sz453_Sx800_Sy800.tif')\n",
-    ")\n",
-    "raw.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "(128, 128, 128)"
-      ]
-     },
-     "execution_count": 6,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "raw = smart_load(\n",
-    "    Path('/Users/qin/Documents/Work/Side_rAP2_16LDs_SAM4_nuclei_Z0.400_X0.291_Y0.291_Sz453_Sx800_Sy800.tif')\n",
-    ")[200:200+128, 300:300+128, 300:300+128]\n",
-    "raw.shape"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[32m2024-12-14 00:35:41.719\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mbioimageio.spec._internal.io_utils\u001b[0m:\u001b[36mopen_bioimageio_yaml\u001b[0m:\u001b[36m131\u001b[0m - \u001b[1mloading emotional-cricket from https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/emotional-cricket/1.1/files/rdf.yaml\u001b[0m\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "INFO: P [MainThread] 2024-12-14 00:35:41,809 plantseg.functionals.prediction.prediction - model expects these inputs: ['raw']\n"
-     ]
-    }
-   ],
-   "source": [
-    "out = biio_prediction(\n",
-    "    raw=raw,\n",
-    "    input_layout='ZYX',\n",
-    "    # model_id='efficient-chipmunk',\n",
-    "    model_id='emotional-cricket',\n",
-    "    # model_id='/Users/qin/Downloads/efficient-chipmunk.yaml',\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 32,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "True"
-      ]
-     },
-     "execution_count": 32,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "len(out.members) == 1"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 28,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['output0']"
-      ]
-     },
-     "execution_count": 28,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "list(out.members)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 31,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[array([[[[[1.23823965e-02, 3.71413166e-03, 2.64202990e-03, ...,\n",
-       "            2.41754763e-03, 4.86477558e-03, 2.36471929e-02],\n",
-       "           [2.31429259e-03, 7.02380203e-04, 4.09954780e-04, ...,\n",
-       "            2.52719648e-04, 7.57241389e-04, 5.46071166e-03],\n",
-       "           [1.33902160e-03, 3.50114511e-04, 2.14061569e-04, ...,\n",
-       "            7.40970063e-05, 2.46272044e-04, 2.46915710e-03],\n",
-       "           ...,\n",
-       "           [8.08374316e-04, 7.53455024e-05, 2.44885123e-05, ...,\n",
-       "            8.84886977e-05, 3.21889500e-04, 2.94925272e-03],\n",
-       "           [1.35905389e-03, 1.61183911e-04, 6.91879104e-05, ...,\n",
-       "            2.89193413e-04, 6.93224138e-04, 4.70515573e-03],\n",
-       "           [9.65875201e-03, 1.88575720e-03, 1.10870227e-03, ...,\n",
-       "            2.58681760e-03, 3.97690758e-03, 1.58278439e-02]],\n",
-       " \n",
-       "          [[3.34995706e-03, 7.49400002e-04, 4.44106583e-04, ...,\n",
-       "            4.23803140e-04, 1.16107799e-03, 7.32442131e-03],\n",
-       "           [4.63154254e-04, 1.31924564e-04, 5.75331833e-05, ...,\n",
-       "            2.48063916e-05, 1.31788765e-04, 1.28855614e-03],\n",
-       "           [1.84829667e-04, 4.22494304e-05, 1.93375345e-05, ...,\n",
-       "            3.12912471e-06, 2.20445181e-05, 3.40537110e-04],\n",
-       "           ...,\n",
-       "           [1.03433682e-04, 4.99963107e-06, 7.01059491e-07, ...,\n",
-       "            5.13319992e-06, 4.15788127e-05, 5.11675724e-04],\n",
-       "           [2.40331719e-04, 2.02179817e-05, 4.81432517e-06, ...,\n",
-       "            2.36316791e-05, 1.03607512e-04, 9.28838330e-04],\n",
-       "           [2.54508085e-03, 3.50280927e-04, 1.46785387e-04, ...,\n",
-       "            3.32942145e-04, 8.13076796e-04, 4.12234711e-03]],\n",
-       " \n",
-       "          [[1.35694118e-03, 1.68786763e-04, 7.48198290e-05, ...,\n",
-       "            2.81019224e-04, 7.46771577e-04, 5.51791256e-03],\n",
-       "           [1.08341432e-04, 1.00760262e-05, 2.75844582e-06, ...,\n",
-       "            8.99239421e-06, 5.42080925e-05, 7.20382726e-04],\n",
-       "           [3.48077228e-05, 2.02582328e-06, 5.01538182e-07, ...,\n",
-       "            5.80834751e-07, 4.99987436e-06, 1.27662483e-04],\n",
-       "           ...,\n",
-       "           [1.09947527e-04, 5.69427766e-06, 8.12078270e-07, ...,\n",
-       "            1.42909403e-06, 1.42246154e-05, 2.88977870e-04],\n",
-       "           [2.98362691e-04, 3.16787191e-05, 7.62458285e-06, ...,\n",
-       "            5.10403333e-06, 3.06438706e-05, 4.58879425e-04],\n",
-       "           [2.99305934e-03, 5.03332238e-04, 2.14670028e-04, ...,\n",
-       "            1.38354051e-04, 4.05884377e-04, 2.70675565e-03]],\n",
-       " \n",
-       "          ...,\n",
-       " \n",
-       "          [[3.29867308e-03, 4.56455105e-04, 1.34146132e-04, ...,\n",
-       "            9.69568064e-05, 3.37075238e-04, 2.77165999e-03],\n",
-       "           [4.24169732e-04, 3.51630406e-05, 5.06967717e-06, ...,\n",
-       "            5.14997964e-06, 3.60815611e-05, 4.70222265e-04],\n",
-       "           [9.61356855e-05, 5.18314391e-06, 7.80053995e-07, ...,\n",
-       "            2.31303488e-06, 1.90853134e-05, 3.29483621e-04],\n",
-       "           ...,\n",
-       "           [8.78510036e-06, 3.06934055e-07, 1.15704395e-07, ...,\n",
-       "            4.02947308e-07, 3.91793219e-06, 1.21722005e-04],\n",
-       "           [3.12283155e-05, 1.88902447e-06, 9.24211236e-07, ...,\n",
-       "            1.67769019e-06, 1.39318608e-05, 3.06272268e-04],\n",
-       "           [9.07270878e-04, 9.94062648e-05, 5.61416928e-05, ...,\n",
-       "            5.73289071e-05, 2.48869706e-04, 2.29152758e-03]],\n",
-       " \n",
-       "          [[5.91580058e-03, 1.10559561e-03, 3.32048250e-04, ...,\n",
-       "            2.32103208e-04, 6.34685624e-04, 4.07362822e-03],\n",
-       "           [9.84633924e-04, 1.30795612e-04, 2.20851525e-05, ...,\n",
-       "            1.67520193e-05, 1.00321871e-04, 8.62589630e-04],\n",
-       "           [2.57799606e-04, 2.33190221e-05, 3.75139280e-06, ...,\n",
-       "            7.50424260e-06, 5.11594008e-05, 5.78730425e-04],\n",
-       "           ...,\n",
-       "           [1.69782816e-05, 6.66840947e-07, 1.38846460e-07, ...,\n",
-       "            1.80018546e-07, 2.14967281e-06, 9.65387953e-05],\n",
-       "           [5.67699681e-05, 3.36161565e-06, 1.02050706e-06, ...,\n",
-       "            1.21392122e-06, 1.14076483e-05, 2.88846204e-04],\n",
-       "           [1.40678498e-03, 1.51812113e-04, 6.07411748e-05, ...,\n",
-       "            5.12399092e-05, 2.33667044e-04, 2.21933913e-03]],\n",
-       " \n",
-       "          [[2.29063556e-02, 6.30153902e-03, 2.76722992e-03, ...,\n",
-       "            2.20514974e-03, 4.18286538e-03, 1.75831020e-02],\n",
-       "           [4.68514580e-03, 7.92992418e-04, 3.06714675e-04, ...,\n",
-       "            2.61121342e-04, 7.31239910e-04, 4.39845258e-03],\n",
-       "           [1.91156555e-03, 2.52487953e-04, 1.00966608e-04, ...,\n",
-       "            1.51911445e-04, 4.64357348e-04, 3.32914991e-03],\n",
-       "           ...,\n",
-       "           [7.02249992e-04, 5.80488013e-05, 1.94775675e-05, ...,\n",
-       "            8.09186895e-06, 4.56816706e-05, 9.44769825e-04],\n",
-       "           [1.51418359e-03, 1.61546574e-04, 7.56840236e-05, ...,\n",
-       "            3.58892830e-05, 1.60053183e-04, 2.05396442e-03],\n",
-       "           [1.25697535e-02, 2.19693198e-03, 1.26698730e-03, ...,\n",
-       "            6.89297158e-04, 1.90667005e-03, 1.04174931e-02]]]]],\n",
-       "       dtype=float32)]"
-      ]
-     },
-     "execution_count": 31,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "[tensor.data.to_numpy() for tensor in out.members.values()]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "array([[[[[1.23823965e-02, 3.71413166e-03, 2.64202990e-03, ...,\n",
-       "           2.41754763e-03, 4.86477558e-03, 2.36471929e-02],\n",
-       "          [2.31429259e-03, 7.02380203e-04, 4.09954780e-04, ...,\n",
-       "           2.52719648e-04, 7.57241389e-04, 5.46071166e-03],\n",
-       "          [1.33902160e-03, 3.50114511e-04, 2.14061569e-04, ...,\n",
-       "           7.40970063e-05, 2.46272044e-04, 2.46915710e-03],\n",
-       "          ...,\n",
-       "          [8.08374316e-04, 7.53455024e-05, 2.44885123e-05, ...,\n",
-       "           8.84886977e-05, 3.21889500e-04, 2.94925272e-03],\n",
-       "          [1.35905389e-03, 1.61183911e-04, 6.91879104e-05, ...,\n",
-       "           2.89193413e-04, 6.93224138e-04, 4.70515573e-03],\n",
-       "          [9.65875201e-03, 1.88575720e-03, 1.10870227e-03, ...,\n",
-       "           2.58681760e-03, 3.97690758e-03, 1.58278439e-02]],\n",
-       "\n",
-       "         [[3.34995706e-03, 7.49400002e-04, 4.44106583e-04, ...,\n",
-       "           4.23803140e-04, 1.16107799e-03, 7.32442131e-03],\n",
-       "          [4.63154254e-04, 1.31924564e-04, 5.75331833e-05, ...,\n",
-       "           2.48063916e-05, 1.31788765e-04, 1.28855614e-03],\n",
-       "          [1.84829667e-04, 4.22494304e-05, 1.93375345e-05, ...,\n",
-       "           3.12912471e-06, 2.20445181e-05, 3.40537110e-04],\n",
-       "          ...,\n",
-       "          [1.03433682e-04, 4.99963107e-06, 7.01059491e-07, ...,\n",
-       "           5.13319992e-06, 4.15788127e-05, 5.11675724e-04],\n",
-       "          [2.40331719e-04, 2.02179817e-05, 4.81432517e-06, ...,\n",
-       "           2.36316791e-05, 1.03607512e-04, 9.28838330e-04],\n",
-       "          [2.54508085e-03, 3.50280927e-04, 1.46785387e-04, ...,\n",
-       "           3.32942145e-04, 8.13076796e-04, 4.12234711e-03]],\n",
-       "\n",
-       "         [[1.35694118e-03, 1.68786763e-04, 7.48198290e-05, ...,\n",
-       "           2.81019224e-04, 7.46771577e-04, 5.51791256e-03],\n",
-       "          [1.08341432e-04, 1.00760262e-05, 2.75844582e-06, ...,\n",
-       "           8.99239421e-06, 5.42080925e-05, 7.20382726e-04],\n",
-       "          [3.48077228e-05, 2.02582328e-06, 5.01538182e-07, ...,\n",
-       "           5.80834751e-07, 4.99987436e-06, 1.27662483e-04],\n",
-       "          ...,\n",
-       "          [1.09947527e-04, 5.69427766e-06, 8.12078270e-07, ...,\n",
-       "           1.42909403e-06, 1.42246154e-05, 2.88977870e-04],\n",
-       "          [2.98362691e-04, 3.16787191e-05, 7.62458285e-06, ...,\n",
-       "           5.10403333e-06, 3.06438706e-05, 4.58879425e-04],\n",
-       "          [2.99305934e-03, 5.03332238e-04, 2.14670028e-04, ...,\n",
-       "           1.38354051e-04, 4.05884377e-04, 2.70675565e-03]],\n",
-       "\n",
-       "         ...,\n",
-       "\n",
-       "         [[3.29867308e-03, 4.56455105e-04, 1.34146132e-04, ...,\n",
-       "           9.69568064e-05, 3.37075238e-04, 2.77165999e-03],\n",
-       "          [4.24169732e-04, 3.51630406e-05, 5.06967717e-06, ...,\n",
-       "           5.14997964e-06, 3.60815611e-05, 4.70222265e-04],\n",
-       "          [9.61356855e-05, 5.18314391e-06, 7.80053995e-07, ...,\n",
-       "           2.31303488e-06, 1.90853134e-05, 3.29483621e-04],\n",
-       "          ...,\n",
-       "          [8.78510036e-06, 3.06934055e-07, 1.15704395e-07, ...,\n",
-       "           4.02947308e-07, 3.91793219e-06, 1.21722005e-04],\n",
-       "          [3.12283155e-05, 1.88902447e-06, 9.24211236e-07, ...,\n",
-       "           1.67769019e-06, 1.39318608e-05, 3.06272268e-04],\n",
-       "          [9.07270878e-04, 9.94062648e-05, 5.61416928e-05, ...,\n",
-       "           5.73289071e-05, 2.48869706e-04, 2.29152758e-03]],\n",
-       "\n",
-       "         [[5.91580058e-03, 1.10559561e-03, 3.32048250e-04, ...,\n",
-       "           2.32103208e-04, 6.34685624e-04, 4.07362822e-03],\n",
-       "          [9.84633924e-04, 1.30795612e-04, 2.20851525e-05, ...,\n",
-       "           1.67520193e-05, 1.00321871e-04, 8.62589630e-04],\n",
-       "          [2.57799606e-04, 2.33190221e-05, 3.75139280e-06, ...,\n",
-       "           7.50424260e-06, 5.11594008e-05, 5.78730425e-04],\n",
-       "          ...,\n",
-       "          [1.69782816e-05, 6.66840947e-07, 1.38846460e-07, ...,\n",
-       "           1.80018546e-07, 2.14967281e-06, 9.65387953e-05],\n",
-       "          [5.67699681e-05, 3.36161565e-06, 1.02050706e-06, ...,\n",
-       "           1.21392122e-06, 1.14076483e-05, 2.88846204e-04],\n",
-       "          [1.40678498e-03, 1.51812113e-04, 6.07411748e-05, ...,\n",
-       "           5.12399092e-05, 2.33667044e-04, 2.21933913e-03]],\n",
-       "\n",
-       "         [[2.29063556e-02, 6.30153902e-03, 2.76722992e-03, ...,\n",
-       "           2.20514974e-03, 4.18286538e-03, 1.75831020e-02],\n",
-       "          [4.68514580e-03, 7.92992418e-04, 3.06714675e-04, ...,\n",
-       "           2.61121342e-04, 7.31239910e-04, 4.39845258e-03],\n",
-       "          [1.91156555e-03, 2.52487953e-04, 1.00966608e-04, ...,\n",
-       "           1.51911445e-04, 4.64357348e-04, 3.32914991e-03],\n",
-       "          ...,\n",
-       "          [7.02249992e-04, 5.80488013e-05, 1.94775675e-05, ...,\n",
-       "           8.09186895e-06, 4.56816706e-05, 9.44769825e-04],\n",
-       "          [1.51418359e-03, 1.61546574e-04, 7.56840236e-05, ...,\n",
-       "           3.58892830e-05, 1.60053183e-04, 2.05396442e-03],\n",
-       "          [1.25697535e-02, 2.19693198e-03, 1.26698730e-03, ...,\n",
-       "           6.89297158e-04, 1.90667005e-03, 1.04174931e-02]]]]],\n",
-       "      dtype=float32)"
-      ]
-     },
-     "execution_count": 24,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out.members['output0'].data.to_numpy()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'output0': <bioimageio.core.tensor.Tensor at 0x3089bf740>}"
-      ]
-     },
-     "execution_count": 17,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out.members."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 34,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "array([[[ 2,  3,  0, ...,  1,  3,  0],\n",
-       "        [ 0,  0,  0, ...,  0,  1,  0],\n",
-       "        [ 3,  3,  2, ...,  0,  1,  0],\n",
-       "        ...,\n",
-       "        [ 3,  4,  2, ...,  1,  2,  1],\n",
-       "        [ 4,  1,  2, ...,  6,  5,  3],\n",
-       "        [ 3,  2,  2, ...,  4,  5,  1]],\n",
-       "\n",
-       "       [[ 2,  2,  4, ...,  5,  8,  9],\n",
-       "        [ 2,  1,  6, ...,  1,  0,  2],\n",
-       "        [ 5,  3,  0, ...,  3,  6,  5],\n",
-       "        ...,\n",
-       "        [ 1,  5,  8, ...,  2,  1,  1],\n",
-       "        [ 3,  3,  6, ...,  2,  3,  3],\n",
-       "        [ 0,  0,  0, ...,  3,  2,  2]],\n",
-       "\n",
-       "       [[ 1,  5,  3, ...,  1,  3,  1],\n",
-       "        [ 5,  3,  8, ...,  2,  5,  6],\n",
-       "        [ 1,  2,  5, ...,  2,  1,  3],\n",
-       "        ...,\n",
-       "        [ 0,  0,  3, ...,  3,  2,  3],\n",
-       "        [ 5,  2,  2, ...,  2,  2,  1],\n",
-       "        [ 6,  2,  3, ...,  8,  2,  6]],\n",
-       "\n",
-       "       ...,\n",
-       "\n",
-       "       [[ 3,  7,  6, ...,  2,  9,  7],\n",
-       "        [ 5,  9,  8, ...,  2,  2,  5],\n",
-       "        [ 4,  1,  2, ...,  6,  6,  3],\n",
-       "        ...,\n",
-       "        [ 0,  3,  6, ...,  2,  3,  3],\n",
-       "        [ 2,  1,  3, ...,  0,  3,  2],\n",
-       "        [ 5,  5,  1, ...,  4,  1,  5]],\n",
-       "\n",
-       "       [[ 9,  6,  5, ...,  2,  3,  3],\n",
-       "        [ 6,  4,  3, ...,  4,  2,  0],\n",
-       "        [ 3,  4,  6, ...,  3,  5,  3],\n",
-       "        ...,\n",
-       "        [ 4,  2,  4, ...,  2,  6,  4],\n",
-       "        [ 0,  0,  4, ...,  0,  2,  2],\n",
-       "        [26,  0,  0, ...,  2,  2,  2]],\n",
-       "\n",
-       "       [[ 4,  2,  4, ...,  0,  1,  1],\n",
-       "        [ 1,  1,  1, ...,  0,  2,  1],\n",
-       "        [ 2,  3,  4, ...,  6,  1,  2],\n",
-       "        ...,\n",
-       "        [ 2,  1,  2, ...,  2,  1,  2],\n",
-       "        [ 4,  6,  3, ...,  1,  1,  2],\n",
-       "        [ 1,  1,  2, ...,  1,  0,  0]]], dtype=uint16)"
-      ]
-     },
-     "execution_count": 34,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "raw"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from plantseg.io import create_tiff\n",
-    "from plantseg.io.voxelsize import VoxelSize"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "create_tiff(Path('/Users/qin/Documents/Work/small_3D_crop.tif'), raw, VoxelSize(voxels_size=(1, 1, 1)))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 42,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "VoxelSize(voxels_size=None, unit='um')"
-      ]
-     },
-     "execution_count": 42,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "VoxelSize()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from bioimageio.spec import load_model_description\n",
-    "from bioimageio.core.axis import AxisId"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[32m2024-12-16 14:36:24.489\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mbioimageio.spec._internal.io_utils\u001b[0m:\u001b[36mopen_bioimageio_yaml\u001b[0m:\u001b[36m131\u001b[0m - \u001b[1mloading philosophical-panda from https://uk1s3.embassy.ebi.ac.uk/public-datasets/bioimage.io/philosophical-panda/0.0.11/files/rdf.yaml\u001b[0m\n",
-      "computing SHA256 of 101853864f8c8e986b2819c9ac44d0f9-test_input.npy (result: 6810255f5b5260fe39153f2192bedf30d9899ec4e770976b7813116c467579f0): 100%|██████████| 5529728/5529728 [00:00<00:00, 1344776498.48it/s]\n",
-      "computing SHA256 of 57b964a123db3ff8400bcce3ef902b18-test_output.npy (result: d802e3024da80bff93a9ec50fbe50b9c3946534aab1b60b911511111a8e2dbca): 100%|██████████| 8294528/8294528 [00:00<00:00, 1573628187.47it/s]\n",
-      "computing SHA256 of 0b712fa8abf8707021a71747e726bf7c-test_style.npy (result: ab464b406f9050561b40f7d76700ab5edf3aca97e31fe9a6069a51aeeca8bc81): 100%|██████████| 76928/76928 [00:00<00:00, 213681733.85it/s]\n",
-      "computing SHA256 of d0abd68ef3844b0d6fbed811e0fed878-test_downsampled_0.npy (result: 67df53fb440e94dbb9c8e4003dcbde158646a7975c4878cacdd251e1fcfb4225): 100%|██████████| 88473728/88473728 [00:00<00:00, 2695492167.77it/s]\n",
-      "computing SHA256 of 3765cac1d92a49daf0d6ec949919aeb1-test_downsampled_1.npy (result: cb4addbd763d96731ebd18ed001b87ab7195ec9198f01a753a363a06c27bfb1c): 100%|██████████| 44236928/44236928 [00:00<00:00, 2611701702.60it/s]\n",
-      "computing SHA256 of 7251078a2afa8713384a3103878dd09d-test_downsampled_2.npy (result: 9c0225b94d84fcc3adfb9a73eef1303d6adb318b57a5a801e0e2e1638b458e72): 100%|██████████| 22118528/22118528 [00:00<00:00, 2409657934.14it/s]\n",
-      "computing SHA256 of e434cdc3ea3e7ecfb752cdc001617875-test_downsampled_3.npy (result: 1ea789ff37d47197c847b585799f7d063e7592b0c5e9c3094fd0e3ac209b7fc2): 100%|██████████| 11059328/11059328 [00:00<00:00, 2449888225.82it/s]\n",
-      "computing SHA256 of 00bd170d6c9de6a391d6869f59058847-cpnet_wrapper.py (result: b8b947cdd0ea8f5b98bd7be5f12f38bb1ea1ebe0b455c62d9a6389cd21d134bf): 100%|██████████| 11053/11053 [00:00<00:00, 34315057.08it/s]\n",
-      "computing SHA256 of 8dbb20d5a3cb3a3dfdb5101a671861ce-cp_state_dict_1135_gold.pth (result: 26c277f3b8f6ca5aab30b4b0a832601aea60183cbed1c2333576f4135a643eb2): 100%|██████████| 26556687/26556687 [00:00<00:00, 2154733982.88it/s]\n",
-      "computing SHA256 of 17fee110c39ccad7c3cb36d00d2fdd2c-cp_traced_1135_gold.pt (result: f61bae146ab522902350eadda1d509ac1037726fe6d7fb63f6a8a314021d63e7): 100%|██████████| 26812339/26812339 [00:00<00:00, 2652775239.24it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "model_id = 'philosophical-panda'\n",
-    "# model_id = 'emotional-cricket'\n",
-    "model = load_model_description(model_id)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/bioimageio/spec/_internal/io.py:351: UserWarning: dumping with mode='python' is currently not fully supported for fields that are included when packaging; returned objects are standard python objects\n",
-      "  warnings.warn(\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=2.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=2.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=2.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=2.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=2.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=2.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=4.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=4.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=4.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=4.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=4.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=4.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=8.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=8.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n",
-      "/Users/qin/micromamba/envs/plant-seg-dev/lib/python3.12/site-packages/pydantic/_internal/_serializers.py:42: UserWarning: Pydantic serializer warnings:\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=1.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=8.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=8.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxis` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=8.0)` - serialized value may not be as expected\n",
-      "  PydanticSerializationUnexpectedValue: Expected `TimeOutputAxisWithHalo` but got `SpaceOutputAxis` with value `SpaceOutputAxis(size=Size...', unit=None, scale=8.0)` - serialized value may not be as expected\n",
-      "  v = handler(item, index)\n"
-     ]
-    },
-    {
-     "data": {
-      "text/plain": [
-       "{'name': 'Cellpose Plant Nuclei ResNet',\n",
-       " 'description': 'An experimental Cellpose nuclear model fine-tuned on ovules 1136, 1137, 1139, 1170 and tested on ovules 1135 (see reference for dataset details). A model for BioImage.IO team to test and develop post-processing tools.',\n",
-       " 'covers': [PosixPath('cellpose_raw_and_segmentation.jpg'),\n",
-       "  PosixPath('cellpose_raw_and_probability.jpg'),\n",
-       "  PosixPath('cellpose_raw.jpg')],\n",
-       " 'id_emoji': '🐼',\n",
-       " 'authors': [{'affiliation': 'EMBL',\n",
-       "   'email': None,\n",
-       "   'orcid': '0000-0002-4652-0795',\n",
-       "   'name': 'Qin Yu',\n",
-       "   'github_user': 'qin-yu'}],\n",
-       " 'attachments': [],\n",
-       " 'cite': [{'text': 'For more details of the model itself, see the manuscript',\n",
-       "   'doi': '10.1101/2024.02.19.580954',\n",
-       "   'url': None}],\n",
-       " 'license': 'MIT',\n",
-       " 'config': {'bioimageio': {'thumbnails': {'cellpose_raw.jpg': 'cellpose_raw.thumbnail.png',\n",
-       "    'cellpose_raw_and_probability.jpg': 'cellpose_raw_and_probability.thumbnail.png',\n",
-       "    'cellpose_raw_and_segmentation.jpg': 'cellpose_raw_and_segmentation.thumbnail.png'}}},\n",
-       " 'git_repo': 'https://github.com/kreshuklab/go-nuclear',\n",
-       " 'icon': None,\n",
-       " 'links': [],\n",
-       " 'uploader': {'email': 'qin.yu.95@outlook.com', 'name': 'Qin Yu'},\n",
-       " 'maintainers': [],\n",
-       " 'tags': ['cellpose', '3d', '2d', 'nuclei'],\n",
-       " 'version': '0.0.11',\n",
-       " 'format_version': '0.5.3',\n",
-       " 'type': 'model',\n",
-       " 'id': 'philosophical-panda',\n",
-       " 'documentation': PosixPath('README.md'),\n",
-       " 'inputs': [{'id': 'raw',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'min': 1, 'step': 1},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0,\n",
-       "     'concatenable': False},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['c1', 'c2']},\n",
-       "    {'size': {'min': 16, 'step': 16},\n",
-       "     'id': 'y',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0,\n",
-       "     'concatenable': False},\n",
-       "    {'size': {'min': 16, 'step': 16},\n",
-       "     'id': 'x',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0,\n",
-       "     'concatenable': False}],\n",
-       "   'test_tensor': {'source': PosixPath('test_input.npy'),\n",
-       "    'sha256': '6810255f5b5260fe39153f2192bedf30d9899ec4e770976b7813116c467579f0'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'optional': False,\n",
-       "   'preprocessing': [{'id': 'ensure_dtype', 'kwargs': {'dtype': 'float32'}}]}],\n",
-       " 'outputs': [{'id': 'flow',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'tensor_id': 'raw', 'axis_id': 'z', 'offset': 0},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['flow1', 'flow2', 'flow3']},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'y', 'offset': 0},\n",
-       "     'id': 'y',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'x', 'offset': 0},\n",
-       "     'id': 'x',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0}],\n",
-       "   'test_tensor': {'source': PosixPath('test_output.npy'),\n",
-       "    'sha256': 'd802e3024da80bff93a9ec50fbe50b9c3946534aab1b60b911511111a8e2dbca'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'postprocessing': [{'id': 'ensure_dtype', 'kwargs': {'dtype': 'float32'}}]},\n",
-       "  {'id': 'style',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'tensor_id': 'raw', 'axis_id': 'z', 'offset': 0},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['feature1',\n",
-       "      'feature2',\n",
-       "      'feature3',\n",
-       "      'feature4',\n",
-       "      'feature5',\n",
-       "      'feature6',\n",
-       "      'feature7',\n",
-       "      'feature8',\n",
-       "      'feature9',\n",
-       "      'feature10',\n",
-       "      'feature11',\n",
-       "      'feature12',\n",
-       "      'feature13',\n",
-       "      'feature14',\n",
-       "      'feature15',\n",
-       "      'feature16',\n",
-       "      'feature17',\n",
-       "      'feature18',\n",
-       "      'feature19',\n",
-       "      'feature20',\n",
-       "      'feature21',\n",
-       "      'feature22',\n",
-       "      'feature23',\n",
-       "      'feature24',\n",
-       "      'feature25',\n",
-       "      'feature26',\n",
-       "      'feature27',\n",
-       "      'feature28',\n",
-       "      'feature29',\n",
-       "      'feature30',\n",
-       "      'feature31',\n",
-       "      'feature32',\n",
-       "      'feature33',\n",
-       "      'feature34',\n",
-       "      'feature35',\n",
-       "      'feature36',\n",
-       "      'feature37',\n",
-       "      'feature38',\n",
-       "      'feature39',\n",
-       "      'feature40',\n",
-       "      'feature41',\n",
-       "      'feature42',\n",
-       "      'feature43',\n",
-       "      'feature44',\n",
-       "      'feature45',\n",
-       "      'feature46',\n",
-       "      'feature47',\n",
-       "      'feature48',\n",
-       "      'feature49',\n",
-       "      'feature50',\n",
-       "      'feature51',\n",
-       "      'feature52',\n",
-       "      'feature53',\n",
-       "      'feature54',\n",
-       "      'feature55',\n",
-       "      'feature56',\n",
-       "      'feature57',\n",
-       "      'feature58',\n",
-       "      'feature59',\n",
-       "      'feature60',\n",
-       "      'feature61',\n",
-       "      'feature62',\n",
-       "      'feature63',\n",
-       "      'feature64',\n",
-       "      'feature65',\n",
-       "      'feature66',\n",
-       "      'feature67',\n",
-       "      'feature68',\n",
-       "      'feature69',\n",
-       "      'feature70',\n",
-       "      'feature71',\n",
-       "      'feature72',\n",
-       "      'feature73',\n",
-       "      'feature74',\n",
-       "      'feature75',\n",
-       "      'feature76',\n",
-       "      'feature77',\n",
-       "      'feature78',\n",
-       "      'feature79',\n",
-       "      'feature80',\n",
-       "      'feature81',\n",
-       "      'feature82',\n",
-       "      'feature83',\n",
-       "      'feature84',\n",
-       "      'feature85',\n",
-       "      'feature86',\n",
-       "      'feature87',\n",
-       "      'feature88',\n",
-       "      'feature89',\n",
-       "      'feature90',\n",
-       "      'feature91',\n",
-       "      'feature92',\n",
-       "      'feature93',\n",
-       "      'feature94',\n",
-       "      'feature95',\n",
-       "      'feature96',\n",
-       "      'feature97',\n",
-       "      'feature98',\n",
-       "      'feature99',\n",
-       "      'feature100',\n",
-       "      'feature101',\n",
-       "      'feature102',\n",
-       "      'feature103',\n",
-       "      'feature104',\n",
-       "      'feature105',\n",
-       "      'feature106',\n",
-       "      'feature107',\n",
-       "      'feature108',\n",
-       "      'feature109',\n",
-       "      'feature110',\n",
-       "      'feature111',\n",
-       "      'feature112',\n",
-       "      'feature113',\n",
-       "      'feature114',\n",
-       "      'feature115',\n",
-       "      'feature116',\n",
-       "      'feature117',\n",
-       "      'feature118',\n",
-       "      'feature119',\n",
-       "      'feature120',\n",
-       "      'feature121',\n",
-       "      'feature122',\n",
-       "      'feature123',\n",
-       "      'feature124',\n",
-       "      'feature125',\n",
-       "      'feature126',\n",
-       "      'feature127',\n",
-       "      'feature128',\n",
-       "      'feature129',\n",
-       "      'feature130',\n",
-       "      'feature131',\n",
-       "      'feature132',\n",
-       "      'feature133',\n",
-       "      'feature134',\n",
-       "      'feature135',\n",
-       "      'feature136',\n",
-       "      'feature137',\n",
-       "      'feature138',\n",
-       "      'feature139',\n",
-       "      'feature140',\n",
-       "      'feature141',\n",
-       "      'feature142',\n",
-       "      'feature143',\n",
-       "      'feature144',\n",
-       "      'feature145',\n",
-       "      'feature146',\n",
-       "      'feature147',\n",
-       "      'feature148',\n",
-       "      'feature149',\n",
-       "      'feature150',\n",
-       "      'feature151',\n",
-       "      'feature152',\n",
-       "      'feature153',\n",
-       "      'feature154',\n",
-       "      'feature155',\n",
-       "      'feature156',\n",
-       "      'feature157',\n",
-       "      'feature158',\n",
-       "      'feature159',\n",
-       "      'feature160',\n",
-       "      'feature161',\n",
-       "      'feature162',\n",
-       "      'feature163',\n",
-       "      'feature164',\n",
-       "      'feature165',\n",
-       "      'feature166',\n",
-       "      'feature167',\n",
-       "      'feature168',\n",
-       "      'feature169',\n",
-       "      'feature170',\n",
-       "      'feature171',\n",
-       "      'feature172',\n",
-       "      'feature173',\n",
-       "      'feature174',\n",
-       "      'feature175',\n",
-       "      'feature176',\n",
-       "      'feature177',\n",
-       "      'feature178',\n",
-       "      'feature179',\n",
-       "      'feature180',\n",
-       "      'feature181',\n",
-       "      'feature182',\n",
-       "      'feature183',\n",
-       "      'feature184',\n",
-       "      'feature185',\n",
-       "      'feature186',\n",
-       "      'feature187',\n",
-       "      'feature188',\n",
-       "      'feature189',\n",
-       "      'feature190',\n",
-       "      'feature191',\n",
-       "      'feature192',\n",
-       "      'feature193',\n",
-       "      'feature194',\n",
-       "      'feature195',\n",
-       "      'feature196',\n",
-       "      'feature197',\n",
-       "      'feature198',\n",
-       "      'feature199',\n",
-       "      'feature200',\n",
-       "      'feature201',\n",
-       "      'feature202',\n",
-       "      'feature203',\n",
-       "      'feature204',\n",
-       "      'feature205',\n",
-       "      'feature206',\n",
-       "      'feature207',\n",
-       "      'feature208',\n",
-       "      'feature209',\n",
-       "      'feature210',\n",
-       "      'feature211',\n",
-       "      'feature212',\n",
-       "      'feature213',\n",
-       "      'feature214',\n",
-       "      'feature215',\n",
-       "      'feature216',\n",
-       "      'feature217',\n",
-       "      'feature218',\n",
-       "      'feature219',\n",
-       "      'feature220',\n",
-       "      'feature221',\n",
-       "      'feature222',\n",
-       "      'feature223',\n",
-       "      'feature224',\n",
-       "      'feature225',\n",
-       "      'feature226',\n",
-       "      'feature227',\n",
-       "      'feature228',\n",
-       "      'feature229',\n",
-       "      'feature230',\n",
-       "      'feature231',\n",
-       "      'feature232',\n",
-       "      'feature233',\n",
-       "      'feature234',\n",
-       "      'feature235',\n",
-       "      'feature236',\n",
-       "      'feature237',\n",
-       "      'feature238',\n",
-       "      'feature239',\n",
-       "      'feature240',\n",
-       "      'feature241',\n",
-       "      'feature242',\n",
-       "      'feature243',\n",
-       "      'feature244',\n",
-       "      'feature245',\n",
-       "      'feature246',\n",
-       "      'feature247',\n",
-       "      'feature248',\n",
-       "      'feature249',\n",
-       "      'feature250',\n",
-       "      'feature251',\n",
-       "      'feature252',\n",
-       "      'feature253',\n",
-       "      'feature254',\n",
-       "      'feature255',\n",
-       "      'feature256']}],\n",
-       "   'test_tensor': {'source': PosixPath('test_style.npy'),\n",
-       "    'sha256': 'ab464b406f9050561b40f7d76700ab5edf3aca97e31fe9a6069a51aeeca8bc81'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'postprocessing': [{'id': 'ensure_dtype', 'kwargs': {'dtype': 'float32'}}]},\n",
-       "  {'id': 'downsampled_0',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'tensor_id': 'raw', 'axis_id': 'z', 'offset': 0},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['feature1',\n",
-       "      'feature2',\n",
-       "      'feature3',\n",
-       "      'feature4',\n",
-       "      'feature5',\n",
-       "      'feature6',\n",
-       "      'feature7',\n",
-       "      'feature8',\n",
-       "      'feature9',\n",
-       "      'feature10',\n",
-       "      'feature11',\n",
-       "      'feature12',\n",
-       "      'feature13',\n",
-       "      'feature14',\n",
-       "      'feature15',\n",
-       "      'feature16',\n",
-       "      'feature17',\n",
-       "      'feature18',\n",
-       "      'feature19',\n",
-       "      'feature20',\n",
-       "      'feature21',\n",
-       "      'feature22',\n",
-       "      'feature23',\n",
-       "      'feature24',\n",
-       "      'feature25',\n",
-       "      'feature26',\n",
-       "      'feature27',\n",
-       "      'feature28',\n",
-       "      'feature29',\n",
-       "      'feature30',\n",
-       "      'feature31',\n",
-       "      'feature32']},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'y', 'offset': 0},\n",
-       "     'id': 'y',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'x', 'offset': 0},\n",
-       "     'id': 'x',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0}],\n",
-       "   'test_tensor': {'source': PosixPath('test_downsampled_0.npy'),\n",
-       "    'sha256': '67df53fb440e94dbb9c8e4003dcbde158646a7975c4878cacdd251e1fcfb4225'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'postprocessing': [{'id': 'ensure_dtype', 'kwargs': {'dtype': 'float32'}}]},\n",
-       "  {'id': 'downsampled_1',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'tensor_id': 'raw', 'axis_id': 'z', 'offset': 0},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['feature1',\n",
-       "      'feature2',\n",
-       "      'feature3',\n",
-       "      'feature4',\n",
-       "      'feature5',\n",
-       "      'feature6',\n",
-       "      'feature7',\n",
-       "      'feature8',\n",
-       "      'feature9',\n",
-       "      'feature10',\n",
-       "      'feature11',\n",
-       "      'feature12',\n",
-       "      'feature13',\n",
-       "      'feature14',\n",
-       "      'feature15',\n",
-       "      'feature16',\n",
-       "      'feature17',\n",
-       "      'feature18',\n",
-       "      'feature19',\n",
-       "      'feature20',\n",
-       "      'feature21',\n",
-       "      'feature22',\n",
-       "      'feature23',\n",
-       "      'feature24',\n",
-       "      'feature25',\n",
-       "      'feature26',\n",
-       "      'feature27',\n",
-       "      'feature28',\n",
-       "      'feature29',\n",
-       "      'feature30',\n",
-       "      'feature31',\n",
-       "      'feature32',\n",
-       "      'feature33',\n",
-       "      'feature34',\n",
-       "      'feature35',\n",
-       "      'feature36',\n",
-       "      'feature37',\n",
-       "      'feature38',\n",
-       "      'feature39',\n",
-       "      'feature40',\n",
-       "      'feature41',\n",
-       "      'feature42',\n",
-       "      'feature43',\n",
-       "      'feature44',\n",
-       "      'feature45',\n",
-       "      'feature46',\n",
-       "      'feature47',\n",
-       "      'feature48',\n",
-       "      'feature49',\n",
-       "      'feature50',\n",
-       "      'feature51',\n",
-       "      'feature52',\n",
-       "      'feature53',\n",
-       "      'feature54',\n",
-       "      'feature55',\n",
-       "      'feature56',\n",
-       "      'feature57',\n",
-       "      'feature58',\n",
-       "      'feature59',\n",
-       "      'feature60',\n",
-       "      'feature61',\n",
-       "      'feature62',\n",
-       "      'feature63',\n",
-       "      'feature64']},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'y', 'offset': 0},\n",
-       "     'id': 'y',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 2.0},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'x', 'offset': 0},\n",
-       "     'id': 'x',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 2.0}],\n",
-       "   'test_tensor': {'source': PosixPath('test_downsampled_1.npy'),\n",
-       "    'sha256': 'cb4addbd763d96731ebd18ed001b87ab7195ec9198f01a753a363a06c27bfb1c'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'postprocessing': [{'id': 'ensure_dtype', 'kwargs': {'dtype': 'float32'}}]},\n",
-       "  {'id': 'downsampled_2',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'tensor_id': 'raw', 'axis_id': 'z', 'offset': 0},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['feature1',\n",
-       "      'feature2',\n",
-       "      'feature3',\n",
-       "      'feature4',\n",
-       "      'feature5',\n",
-       "      'feature6',\n",
-       "      'feature7',\n",
-       "      'feature8',\n",
-       "      'feature9',\n",
-       "      'feature10',\n",
-       "      'feature11',\n",
-       "      'feature12',\n",
-       "      'feature13',\n",
-       "      'feature14',\n",
-       "      'feature15',\n",
-       "      'feature16',\n",
-       "      'feature17',\n",
-       "      'feature18',\n",
-       "      'feature19',\n",
-       "      'feature20',\n",
-       "      'feature21',\n",
-       "      'feature22',\n",
-       "      'feature23',\n",
-       "      'feature24',\n",
-       "      'feature25',\n",
-       "      'feature26',\n",
-       "      'feature27',\n",
-       "      'feature28',\n",
-       "      'feature29',\n",
-       "      'feature30',\n",
-       "      'feature31',\n",
-       "      'feature32',\n",
-       "      'feature33',\n",
-       "      'feature34',\n",
-       "      'feature35',\n",
-       "      'feature36',\n",
-       "      'feature37',\n",
-       "      'feature38',\n",
-       "      'feature39',\n",
-       "      'feature40',\n",
-       "      'feature41',\n",
-       "      'feature42',\n",
-       "      'feature43',\n",
-       "      'feature44',\n",
-       "      'feature45',\n",
-       "      'feature46',\n",
-       "      'feature47',\n",
-       "      'feature48',\n",
-       "      'feature49',\n",
-       "      'feature50',\n",
-       "      'feature51',\n",
-       "      'feature52',\n",
-       "      'feature53',\n",
-       "      'feature54',\n",
-       "      'feature55',\n",
-       "      'feature56',\n",
-       "      'feature57',\n",
-       "      'feature58',\n",
-       "      'feature59',\n",
-       "      'feature60',\n",
-       "      'feature61',\n",
-       "      'feature62',\n",
-       "      'feature63',\n",
-       "      'feature64',\n",
-       "      'feature65',\n",
-       "      'feature66',\n",
-       "      'feature67',\n",
-       "      'feature68',\n",
-       "      'feature69',\n",
-       "      'feature70',\n",
-       "      'feature71',\n",
-       "      'feature72',\n",
-       "      'feature73',\n",
-       "      'feature74',\n",
-       "      'feature75',\n",
-       "      'feature76',\n",
-       "      'feature77',\n",
-       "      'feature78',\n",
-       "      'feature79',\n",
-       "      'feature80',\n",
-       "      'feature81',\n",
-       "      'feature82',\n",
-       "      'feature83',\n",
-       "      'feature84',\n",
-       "      'feature85',\n",
-       "      'feature86',\n",
-       "      'feature87',\n",
-       "      'feature88',\n",
-       "      'feature89',\n",
-       "      'feature90',\n",
-       "      'feature91',\n",
-       "      'feature92',\n",
-       "      'feature93',\n",
-       "      'feature94',\n",
-       "      'feature95',\n",
-       "      'feature96',\n",
-       "      'feature97',\n",
-       "      'feature98',\n",
-       "      'feature99',\n",
-       "      'feature100',\n",
-       "      'feature101',\n",
-       "      'feature102',\n",
-       "      'feature103',\n",
-       "      'feature104',\n",
-       "      'feature105',\n",
-       "      'feature106',\n",
-       "      'feature107',\n",
-       "      'feature108',\n",
-       "      'feature109',\n",
-       "      'feature110',\n",
-       "      'feature111',\n",
-       "      'feature112',\n",
-       "      'feature113',\n",
-       "      'feature114',\n",
-       "      'feature115',\n",
-       "      'feature116',\n",
-       "      'feature117',\n",
-       "      'feature118',\n",
-       "      'feature119',\n",
-       "      'feature120',\n",
-       "      'feature121',\n",
-       "      'feature122',\n",
-       "      'feature123',\n",
-       "      'feature124',\n",
-       "      'feature125',\n",
-       "      'feature126',\n",
-       "      'feature127',\n",
-       "      'feature128']},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'y', 'offset': 0},\n",
-       "     'id': 'y',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 4.0},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'x', 'offset': 0},\n",
-       "     'id': 'x',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 4.0}],\n",
-       "   'test_tensor': {'source': PosixPath('test_downsampled_2.npy'),\n",
-       "    'sha256': '9c0225b94d84fcc3adfb9a73eef1303d6adb318b57a5a801e0e2e1638b458e72'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'postprocessing': [{'id': 'ensure_dtype', 'kwargs': {'dtype': 'float32'}}]},\n",
-       "  {'id': 'downsampled_3',\n",
-       "   'description': '',\n",
-       "   'axes': [{'size': {'tensor_id': 'raw', 'axis_id': 'z', 'offset': 0},\n",
-       "     'id': 'z',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 1.0},\n",
-       "    {'id': 'channel',\n",
-       "     'description': '',\n",
-       "     'type': 'channel',\n",
-       "     'channel_names': ['feature1',\n",
-       "      'feature2',\n",
-       "      'feature3',\n",
-       "      'feature4',\n",
-       "      'feature5',\n",
-       "      'feature6',\n",
-       "      'feature7',\n",
-       "      'feature8',\n",
-       "      'feature9',\n",
-       "      'feature10',\n",
-       "      'feature11',\n",
-       "      'feature12',\n",
-       "      'feature13',\n",
-       "      'feature14',\n",
-       "      'feature15',\n",
-       "      'feature16',\n",
-       "      'feature17',\n",
-       "      'feature18',\n",
-       "      'feature19',\n",
-       "      'feature20',\n",
-       "      'feature21',\n",
-       "      'feature22',\n",
-       "      'feature23',\n",
-       "      'feature24',\n",
-       "      'feature25',\n",
-       "      'feature26',\n",
-       "      'feature27',\n",
-       "      'feature28',\n",
-       "      'feature29',\n",
-       "      'feature30',\n",
-       "      'feature31',\n",
-       "      'feature32',\n",
-       "      'feature33',\n",
-       "      'feature34',\n",
-       "      'feature35',\n",
-       "      'feature36',\n",
-       "      'feature37',\n",
-       "      'feature38',\n",
-       "      'feature39',\n",
-       "      'feature40',\n",
-       "      'feature41',\n",
-       "      'feature42',\n",
-       "      'feature43',\n",
-       "      'feature44',\n",
-       "      'feature45',\n",
-       "      'feature46',\n",
-       "      'feature47',\n",
-       "      'feature48',\n",
-       "      'feature49',\n",
-       "      'feature50',\n",
-       "      'feature51',\n",
-       "      'feature52',\n",
-       "      'feature53',\n",
-       "      'feature54',\n",
-       "      'feature55',\n",
-       "      'feature56',\n",
-       "      'feature57',\n",
-       "      'feature58',\n",
-       "      'feature59',\n",
-       "      'feature60',\n",
-       "      'feature61',\n",
-       "      'feature62',\n",
-       "      'feature63',\n",
-       "      'feature64',\n",
-       "      'feature65',\n",
-       "      'feature66',\n",
-       "      'feature67',\n",
-       "      'feature68',\n",
-       "      'feature69',\n",
-       "      'feature70',\n",
-       "      'feature71',\n",
-       "      'feature72',\n",
-       "      'feature73',\n",
-       "      'feature74',\n",
-       "      'feature75',\n",
-       "      'feature76',\n",
-       "      'feature77',\n",
-       "      'feature78',\n",
-       "      'feature79',\n",
-       "      'feature80',\n",
-       "      'feature81',\n",
-       "      'feature82',\n",
-       "      'feature83',\n",
-       "      'feature84',\n",
-       "      'feature85',\n",
-       "      'feature86',\n",
-       "      'feature87',\n",
-       "      'feature88',\n",
-       "      'feature89',\n",
-       "      'feature90',\n",
-       "      'feature91',\n",
-       "      'feature92',\n",
-       "      'feature93',\n",
-       "      'feature94',\n",
-       "      'feature95',\n",
-       "      'feature96',\n",
-       "      'feature97',\n",
-       "      'feature98',\n",
-       "      'feature99',\n",
-       "      'feature100',\n",
-       "      'feature101',\n",
-       "      'feature102',\n",
-       "      'feature103',\n",
-       "      'feature104',\n",
-       "      'feature105',\n",
-       "      'feature106',\n",
-       "      'feature107',\n",
-       "      'feature108',\n",
-       "      'feature109',\n",
-       "      'feature110',\n",
-       "      'feature111',\n",
-       "      'feature112',\n",
-       "      'feature113',\n",
-       "      'feature114',\n",
-       "      'feature115',\n",
-       "      'feature116',\n",
-       "      'feature117',\n",
-       "      'feature118',\n",
-       "      'feature119',\n",
-       "      'feature120',\n",
-       "      'feature121',\n",
-       "      'feature122',\n",
-       "      'feature123',\n",
-       "      'feature124',\n",
-       "      'feature125',\n",
-       "      'feature126',\n",
-       "      'feature127',\n",
-       "      'feature128',\n",
-       "      'feature129',\n",
-       "      'feature130',\n",
-       "      'feature131',\n",
-       "      'feature132',\n",
-       "      'feature133',\n",
-       "      'feature134',\n",
-       "      'feature135',\n",
-       "      'feature136',\n",
-       "      'feature137',\n",
-       "      'feature138',\n",
-       "      'feature139',\n",
-       "      'feature140',\n",
-       "      'feature141',\n",
-       "      'feature142',\n",
-       "      'feature143',\n",
-       "      'feature144',\n",
-       "      'feature145',\n",
-       "      'feature146',\n",
-       "      'feature147',\n",
-       "      'feature148',\n",
-       "      'feature149',\n",
-       "      'feature150',\n",
-       "      'feature151',\n",
-       "      'feature152',\n",
-       "      'feature153',\n",
-       "      'feature154',\n",
-       "      'feature155',\n",
-       "      'feature156',\n",
-       "      'feature157',\n",
-       "      'feature158',\n",
-       "      'feature159',\n",
-       "      'feature160',\n",
-       "      'feature161',\n",
-       "      'feature162',\n",
-       "      'feature163',\n",
-       "      'feature164',\n",
-       "      'feature165',\n",
-       "      'feature166',\n",
-       "      'feature167',\n",
-       "      'feature168',\n",
-       "      'feature169',\n",
-       "      'feature170',\n",
-       "      'feature171',\n",
-       "      'feature172',\n",
-       "      'feature173',\n",
-       "      'feature174',\n",
-       "      'feature175',\n",
-       "      'feature176',\n",
-       "      'feature177',\n",
-       "      'feature178',\n",
-       "      'feature179',\n",
-       "      'feature180',\n",
-       "      'feature181',\n",
-       "      'feature182',\n",
-       "      'feature183',\n",
-       "      'feature184',\n",
-       "      'feature185',\n",
-       "      'feature186',\n",
-       "      'feature187',\n",
-       "      'feature188',\n",
-       "      'feature189',\n",
-       "      'feature190',\n",
-       "      'feature191',\n",
-       "      'feature192',\n",
-       "      'feature193',\n",
-       "      'feature194',\n",
-       "      'feature195',\n",
-       "      'feature196',\n",
-       "      'feature197',\n",
-       "      'feature198',\n",
-       "      'feature199',\n",
-       "      'feature200',\n",
-       "      'feature201',\n",
-       "      'feature202',\n",
-       "      'feature203',\n",
-       "      'feature204',\n",
-       "      'feature205',\n",
-       "      'feature206',\n",
-       "      'feature207',\n",
-       "      'feature208',\n",
-       "      'feature209',\n",
-       "      'feature210',\n",
-       "      'feature211',\n",
-       "      'feature212',\n",
-       "      'feature213',\n",
-       "      'feature214',\n",
-       "      'feature215',\n",
-       "      'feature216',\n",
-       "      'feature217',\n",
-       "      'feature218',\n",
-       "      'feature219',\n",
-       "      'feature220',\n",
-       "      'feature221',\n",
-       "      'feature222',\n",
-       "      'feature223',\n",
-       "      'feature224',\n",
-       "      'feature225',\n",
-       "      'feature226',\n",
-       "      'feature227',\n",
-       "      'feature228',\n",
-       "      'feature229',\n",
-       "      'feature230',\n",
-       "      'feature231',\n",
-       "      'feature232',\n",
-       "      'feature233',\n",
-       "      'feature234',\n",
-       "      'feature235',\n",
-       "      'feature236',\n",
-       "      'feature237',\n",
-       "      'feature238',\n",
-       "      'feature239',\n",
-       "      'feature240',\n",
-       "      'feature241',\n",
-       "      'feature242',\n",
-       "      'feature243',\n",
-       "      'feature244',\n",
-       "      'feature245',\n",
-       "      'feature246',\n",
-       "      'feature247',\n",
-       "      'feature248',\n",
-       "      'feature249',\n",
-       "      'feature250',\n",
-       "      'feature251',\n",
-       "      'feature252',\n",
-       "      'feature253',\n",
-       "      'feature254',\n",
-       "      'feature255',\n",
-       "      'feature256']},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'y', 'offset': 0},\n",
-       "     'id': 'y',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 8.0},\n",
-       "    {'size': {'tensor_id': 'raw', 'axis_id': 'x', 'offset': 0},\n",
-       "     'id': 'x',\n",
-       "     'description': '',\n",
-       "     'type': 'space',\n",
-       "     'unit': None,\n",
-       "     'scale': 8.0}],\n",
-       "   'test_tensor': {'source': PosixPath('test_downsampled_3.npy'),\n",
-       "    'sha256': '1ea789ff37d47197c847b585799f7d063e7592b0c5e9c3094fd0e3ac209b7fc2'},\n",
-       "   'sample_tensor': None,\n",
-       "   'data': {'type': 'float32',\n",
-       "    'range': (None, None),\n",
-       "    'unit': 'arbitrary unit',\n",
-       "    'scale': 1.0,\n",
-       "    'offset': None},\n",
-       "   'postprocessing': [{'id': 'ensure_dtype',\n",
-       "     'kwargs': {'dtype': 'float32'}}]}],\n",
-       " 'packaged_by': [],\n",
-       " 'parent': None,\n",
-       " 'run_mode': None,\n",
-       " 'timestamp': datetime.datetime(2024, 12, 16, 13, 13, 2, 698812),\n",
-       " 'training_data': None,\n",
-       " 'weights': {'keras_hdf5': None,\n",
-       "  'onnx': None,\n",
-       "  'pytorch_state_dict': {'source': PosixPath('cp_state_dict_1135_gold.pth'),\n",
-       "   'sha256': '26c277f3b8f6ca5aab30b4b0a832601aea60183cbed1c2333576f4135a643eb2',\n",
-       "   'authors': None,\n",
-       "   'parent': None,\n",
-       "   'architecture': {'source': PosixPath('cpnet_wrapper.py'),\n",
-       "    'sha256': 'b8b947cdd0ea8f5b98bd7be5f12f38bb1ea1ebe0b455c62d9a6389cd21d134bf',\n",
-       "    'callable': 'CPnetBioImageIO',\n",
-       "    'kwargs': {'conv_3D': False,\n",
-       "     'max_pool': True,\n",
-       "     'mkldnn': False,\n",
-       "     'nbase': [2, 32, 64, 128, 256],\n",
-       "     'nout': 3,\n",
-       "     'sz': 3}},\n",
-       "   'pytorch_version': '2.3.1',\n",
-       "   'dependencies': None},\n",
-       "  'tensorflow_js': None,\n",
-       "  'tensorflow_saved_model_bundle': None,\n",
-       "  'torchscript': {'source': PosixPath('cp_traced_1135_gold.pt'),\n",
-       "   'sha256': 'f61bae146ab522902350eadda1d509ac1037726fe6d7fb63f6a8a314021d63e7',\n",
-       "   'authors': None,\n",
-       "   'parent': 'pytorch_state_dict',\n",
-       "   'pytorch_version': '2.3.1'}}}"
-      ]
-     },
-     "execution_count": 12,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "model.model_dump()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "['batch', 'channel', 'z', 'y', 'x']"
-      ]
-     },
-     "execution_count": 10,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "[AxisId(a) if isinstance(a, str) else a.id for a in model.inputs[0].axes]"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "{AxisId(a) if isinstance(a, str) else a.id : a.get('size', 1) for a in model.inputs[0].axes}"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "plant-seg-dev",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.12.7"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}