From 78c4325f96e393e5d9a77d88f850920224083a26 Mon Sep 17 00:00:00 2001
From: Jair Sanchez <jair.sanchez@isas.de>
Date: Thu, 13 Nov 2025 17:36:09 +0100
Subject: [PATCH] Preprocessing correction

---
 mmv_im2im/__init__.py      | 2 +-
 mmv_im2im/map_extractor.py | 3 +++
 pyproject.toml             | 2 +-
 setup.cfg                  | 4 ++--
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/mmv_im2im/__init__.py b/mmv_im2im/__init__.py
index a2a3f7e..dd33d80 100644
--- a/mmv_im2im/__init__.py
+++ b/mmv_im2im/__init__.py
@@ -6,7 +6,7 @@
 __email__ = "jianxuchen.ai@gmail.com"
 # Do not edit this string manually, always use bumpversion
 # Details in CONTRIBUTING.md
-__version__ = "0.6.1"
+__version__ = "0.5.7"
 
 
 def get_module_version():
diff --git a/mmv_im2im/map_extractor.py b/mmv_im2im/map_extractor.py
index b84acb0..ee0b903 100644
--- a/mmv_im2im/map_extractor.py
+++ b/mmv_im2im/map_extractor.py
@@ -135,6 +135,7 @@ def process_one_image(
 
         if self.pre_process is not None:
             x = self.pre_process(x)
+            x = x[0]
 
         # choose different inference function for different types of models
         # the input here is assumed to be a tensor
@@ -257,6 +258,8 @@ def process_one_image(
     def run_inference(self):
 
         self.setup_model()
+        self.setup_data_processing()
+
         if "pred_slice2vol" in self.model_cfg.net:
 
             if self.model_cfg.net["pred_slice2vol"] is not None:
diff --git a/pyproject.toml b/pyproject.toml
index f1ff917..3de9e79 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mmv_im2im"
-version = "0.6.1"
+version = "0.5.7"
 authors = [
     { name="Jianxu Chen", email="jianxuchen.ai@gmail.com" },
 ]
diff --git a/setup.cfg b/setup.cfg
index e2ac050..4010c51 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 0.6.1
+current_version = 0.5.7
 commit = True
 tag = True
 
@@ -14,4 +14,4 @@ replace = version = "{new_version}"
 [flake8]
 exclude = docs/, .git/, __pycache__/, build/, dist/, .venv/, .tox/, *.egg-info/
 ignore = E203, E402, W291, W503, W293, W292, E501
-max-line-length = 88
+max-line-length = 88
\ No newline at end of file