diff --git a/aliby/pipeline.py b/aliby/pipeline.py
index 0e175f4ac6378e8d095690a3a88ce6f163a02e6e..7f02704b4d8ed8bbb912a1c4ae2218c3ee4bb575 100644
--- a/aliby/pipeline.py
+++ b/aliby/pipeline.py
@@ -38,7 +38,7 @@ from tqdm import tqdm
 from aliby.baby_client import BabyParameters, BabyRunner
 from aliby.haystack import initialise_tf
 from aliby.io.dataset import Dataset, DatasetLocal
-from aliby.io.image import Image, ImageLocal
+from aliby.io.image import get_image_class
 from aliby.tile.tiler import Tiler, TilerParameters
 from extraction.core.extractor import Extractor, ExtractorParameters
 from extraction.core.functions.defaults import exparams_from_meta
@@ -361,7 +361,7 @@ class Pipeline(ProcessABC):
             # Set up
             directory = general_config["directory"]
 
-            image_wrapper = ImageLocal if isinstance(image_id, str) else Image
+            image_wrapper = get_image_class(image_id)
 
             with image_wrapper(
                 image_id, **self.general.get("server_info", {})