add load_image param

Signed-off-by: binliu <[email protected]>

vista3d_pipeline.py

@@ -52,6 +52,7 @@ class VISTA3DPipeline(Pipeline):
         "image_key",
         "resample_spacing",
         "metadata_path",
+        "load_image",
     ]
     INFERENCE_EXTRA_ARGS = [
         "mode",
@@ -105,6 +106,7 @@ class VISTA3DPipeline(Pipeline):
         image_key: str = "image",
         resample_spacing: Sequence = (1.5, 1.5, 1.5),
         metadata_path: str = os.path.join(FILE_PATH, "metadata.json"),
+        load_image: bool = True,
     ):
         device = self.device
         subclass = {
@@ -114,8 +116,7 @@ class VISTA3DPipeline(Pipeline):
         }
         metadata = json.loads(pathlib.Path(metadata_path).read_text())
         labels_dict = metadata["network_data_format"]["outputs"]["pred"]["channel_def"]
-        preprocessing_transforms = Compose(
-            [
+        preprocessing_list = [
                 LoadImaged(keys=image_key, image_only=True),
                 EnsureChannelFirstd(keys=image_key),
                 EnsureTyped(keys=image_key, device=device, track_meta=True),
@@ -137,7 +138,10 @@ class VISTA3DPipeline(Pipeline):
                 Orientationd(keys=image_key, axcodes="RAS"),
                 CastToTyped(keys=image_key, dtype=torch.float32),
             ]
-        )
+        if not load_image:
+            preprocessing_list.pop(0)
+
+        preprocessing_transforms = Compose(preprocessing_list)
         return preprocessing_transforms
 
     def _init_postprocessing_transforms(