ADD: add multifile and multichannel flag to local datastore

Adds flags to the local datastore init signaling the data is either multichannel or multi-file (each directory has a sample with multiple volumes

Signed-off-by: Cavan Riley <cavan-riley@uiowa.edu>

Author: CavRiley

monailabel/datastore/local.py

@@ -102,9 +102,11 @@ def __init__(
         images_dir: str = ".",
         labels_dir: str = "labels",
         datastore_config: str = "datastore_v2.json",
-        extensions=("*.nii.gz", "*.nii"),
+        extensions=("*.nii.gz", "*.nii", "*.nrrd"),
         auto_reload=False,
         read_only=False,
+        multichannel=False,
+        multi_file=False,
     ):
         """
         Creates a `LocalDataset` object
@@ -124,6 +126,8 @@ def __init__(
         self._ignore_event_config = False
         self._config_ts = 0
         self._auto_reload = auto_reload
+        self._multichannel = multichannel
+        self._multi_file = multi_file
 
         logging.getLogger("filelock").setLevel(logging.ERROR)
 
@@ -256,6 +260,12 @@ def datalist(self, full_path=True) -> List[Dict[str, Any]]:
             ds = json.loads(json.dumps(ds).replace(f"{self._datastore_path.rstrip(os.pathsep)}{os.pathsep}", ""))
         return ds
 
+    def get_is_multichannel(self) -> bool:
+        return self._multichannel
+
+    def get_is_multi_file(self) -> bool:
+        return self._multi_file
+
     def get_image(self, image_id: str, params=None) -> Any:
         """
         Retrieve image object based on image id
@@ -431,6 +441,29 @@ def refresh(self):
         """
         self._reconcile_datastore()
 
+    def add_directory(self, directory_id: str, filename: str, info: Dict[str, Any]) -> str:
+        id = os.path.basename(filename)
+        if not directory_id:
+            directory_id = id
+
+        logger.info(f"Adding Image: {directory_id} => {filename}")
+        name = directory_id
+        dest = os.path.realpath(os.path.join(self._datastore.image_path(), name))
+
+        with FileLock(self._lock_file):
+            logger.debug("Acquired the lock!")
+            shutil.copy(filename, dest)
+
+            info = info if info else {}
+            info["ts"] = int(time.time())
+            info["name"] = name
+
+            # images = get_directory_contents(filename)
+            self._datastore.objects[directory_id] = ImageLabelModel(image=DataModel(info=info, ext=""))
+            self._update_datastore_file(lock=False)
+        logger.debug("Released the lock!")
+        return directory_id
+
     def add_image(self, image_id: str, image_filename: str, image_info: Dict[str, Any]) -> str:
         id, image_ext = self._to_id(os.path.basename(image_filename))
         if not image_id:
@@ -552,10 +585,15 @@ def _list_files(self, path, patterns):
         files = os.listdir(path)
 
         filtered = dict()
-        for pattern in patterns:
-            matching = fnmatch.filter(files, pattern)
-            for file in matching:
-                filtered[os.path.basename(file)] = file
+        if not self._multi_file:
+            for pattern in patterns:
+                matching = fnmatch.filter(files, pattern)
+                for file in matching:
+                    filtered[os.path.basename(file)] = file
+        else:
+            for file in files:
+                if file.lower() not in ["labels", ".lock", "datastore_v2.json"]:
+                    filtered[os.path.basename(file)] = file
         return filtered
 
     def _reconcile_datastore(self):
@@ -585,23 +623,26 @@ def _add_non_existing_images(self) -> int:
         invalidate = 0
         self._init_from_datastore_file()
 
-        local_images = self._list_files(self._datastore.image_path(), self._extensions)
+        local_files = self._list_files(self._datastore.image_path(), self._extensions)
 
-        image_ids = list(self._datastore.objects.keys())
-        for image_file in local_images:
-            image_id, image_ext = self._to_id(image_file)
-            if image_id not in image_ids:
-                logger.info(f"Adding New Image: {image_id} => {image_file}")
+        ids = list(self._datastore.objects.keys())
+        for file in local_files:
+            if self._multi_file:
+                # Directories have no extension — use the name as-is
+                file_id = file
+                file_ext_str = ""
+            else:
+                file_id, file_ext_str = self._to_id(file)
 
-                name = self._filename(image_id, image_ext)
-                image_info = {
+            if file_id not in ids:
+                logger.info(f"Adding New Image: {file_id} => {file}")
+                name = self._filename(file_id, file_ext_str)
+                file_info = {
                     "ts": int(time.time()),
-                    # "checksum": file_checksum(os.path.join(self._datastore.image_path(), name)),
                     "name": name,
                 }
-
                 invalidate += 1
-                self._datastore.objects[image_id] = ImageLabelModel(image=DataModel(info=image_info, ext=image_ext))
+                self._datastore.objects[file_id] = ImageLabelModel(name=DataModel(info=file_info, ext=file_ext_str))
 
         return invalidate
 

monailabel/endpoints/datastore.py

@@ -68,7 +68,7 @@ def add_image(
     logger.info(f"Image: {image}; File: {file}; params: {params}")
     file_ext = "".join(pathlib.Path(file.filename).suffixes) if file.filename else ".nii.gz"
 
-    image_id = image if image else os.path.basename(file.filename).replace(file_ext, "")
+    id = image if image else os.path.basename(file.filename).replace(file_ext, "")
     image_file = tempfile.NamedTemporaryFile(suffix=file_ext).name
 
     with open(image_file, "wb") as buffer:
@@ -79,8 +79,12 @@ def add_image(
     save_params: Dict[str, Any] = json.loads(params) if params else {}
     if user:
         save_params["user"] = user
-    image_id = instance.datastore().add_image(image_id, image_file, save_params)
-    return {"image": image_id}
+    if not instance.datastore().get_is_multi_file():
+        image_id = instance.datastore().add_image(id, image_file, save_params)
+        return {"image": image_id}
+    else:
+        directory_id = instance.datastore().add_directory(id, image_file, save_params)
+        return {"image": directory_id}
 
 
 def remove_image(id: str, user: Optional[str] = None):

monailabel/interfaces/app.py

@@ -90,7 +90,9 @@ def __init__(
         self.app_dir = app_dir
         self.studies = studies
         self.conf = conf if conf else {}
-
+        self.multichannel = conf.get("multichannel", False)
+        self.multi_file = conf.get("multi_file", False)
+        self.input_channels = conf.get("input_channels", False)
         self.name = name
         self.description = description
         self.version = version
@@ -142,6 +144,8 @@ def init_datastore(self) -> Datastore:
             extensions=settings.MONAI_LABEL_DATASTORE_FILE_EXT,
             auto_reload=settings.MONAI_LABEL_DATASTORE_AUTO_RELOAD,
             read_only=settings.MONAI_LABEL_DATASTORE_READ_ONLY,
+            multichannel=self.multichannel,
+            multi_file=self.multi_file,
         )
 
     def init_remote_datastore(self) -> Datastore:
@@ -277,6 +281,10 @@ def infer(self, request, datastore=None):
                 f"Inference Task is not Initialized. There is no model '{model}' available",
             )
 
+        request["multi_file"] = self.multi_file
+        request["multichannel"] = self.multichannel
+        request["input_channels"] = self.input_channels
+
         request = copy.deepcopy(request)
         request["description"] = task.description
 
@@ -288,7 +296,7 @@ def infer(self, request, datastore=None):
             else:
                 request["image"] = datastore.get_image_uri(request["image"])
 
-            if os.path.isdir(request["image"]):
+            if os.path.isdir(request["image"]) and not self.multi_file:
                 logger.info("Input is a Directory; Consider it as DICOM")
 
             logger.debug(f"Image => {request['image']}")
@@ -419,6 +427,11 @@ def train(self, request):
                 f"Train Task is not Initialized. There is no model '{model}' available; {request}",
             )
 
+        # 4D image support, send train task information regarding data
+        request["multi_file"] = self.multi_file
+        request["multichannel"] = self.multichannel
+        request["input_channels"] = self.input_channels
+
         request = copy.deepcopy(request)
         result = task(request, self.datastore())
 

monailabel/interfaces/datastore.py

@@ -201,6 +201,19 @@ def refresh(self) -> None:
         """
         pass
 
+    # TODO: look into implementing for all datastore children and make abstract
+    # @abstractmethod
+    def add_directory(self, id: str, filename: str, info: Dict[str, Any]) -> str:
+        """
+        Save a directory for the given directory id and return the newly saved directory's id
+
+        :param id: the directory id for the image;  If None then base filename will be used
+        :param filename: the path to the directory
+        :param info: additional info for the directory
+        :return: the directory id for the saved image filename
+        """
+        pass
+
     @abstractmethod
     def add_image(self, image_id: str, image_filename: str, image_info: Dict[str, Any]) -> str:
         """
@@ -279,3 +292,19 @@ def json(self):
         Return json representation of datastore
         """
         pass
+
+    # TODO: look into implementing for all datastore children and make abstract
+    # @abstractmethod
+    def get_is_multichannel(self) -> bool:
+        """
+        Returns whether the application's studies is directed at multichannel (4D) data
+        """
+        pass
+
+    # TODO: look into implementing for all datastore children and make abstract
+    # @abstractmethod
+    def get_is_multi_file(self) -> bool:
+        """
+        Returns whether the application's studies is directed at directories containing multiple images per sample
+        """
+        pass

monailabel/tasks/activelearning/first.py

@@ -35,5 +35,13 @@ def __call__(self, request, datastore: Datastore):
         images.sort()
         image = images[0]
 
+        # If the datastore contains 4d images send the multichannel flag to ensure images are loaded as sequences
+        if datastore.get_is_multichannel():
+            return {"id": image, "multichannel": True}
+
+        # If the datastore is multi_file, each sample has a directory with multiple images
+        if datastore.get_is_multi_file():
+            return {"id": image, "multi_file": True}
+
         logger.info(f"First: Selected Image: {image}")
         return {"id": image}

monailabel/tasks/activelearning/random.py

@@ -45,4 +45,17 @@ def __call__(self, request, datastore: Datastore):
         image = random.choices(images, weights=weights)[0]
         logger.debug(f"Random: Images: {images}; Weight: {weights}")
         logger.info(f"Random: Selected Image: {image}; Weight: {weights[0]}")
+
+        # If the datastore contains 4d images send the multichannel flag to ensure images are loaded as sequences
+        if datastore.get_is_multichannel():
+            return {"id": image, "weight": weights[0], "multichannel": True}
+
+        # If the datastore is multi_file, each sample has a directory with multiple images
+        if datastore.get_is_multi_file():
+            return {
+                "id": image,
+                "weight": weights[0],
+                "multi_file": True,
+            }  # this will send the directory and we will walk it later on
+
         return {"id": image, "weight": weights[0]}

monailabel/tasks/train/basic_train.py

@@ -83,6 +83,8 @@ def __init__(self):
         self.multi_gpu = False  # multi gpu enabled
         self.local_rank = 0  # local rank in case of multi gpu
         self.world_size = 0  # world size in case of multi gpu
+        self.input_channels = 1
+        self.multi_file = False
 
         self.request = None
         self.trainer = None
@@ -490,6 +492,9 @@ def train(self, rank, world_size, request, datalist):
 
         context.run_id = request["run_id"]
         context.multi_gpu = request["multi_gpu"]
+        context.multi_file = request.get("multi_file", False)
+        context.input_channels = request.get("input_channels", 1)
+
         if context.multi_gpu:
             os.environ["LOCAL_RANK"] = str(context.local_rank)