Multichannel Model Support for Pre-Aligned Volumes

monailabel/datastore/cvat.py

@@ -16,6 +16,7 @@
 import tempfile
 import time
 import urllib.parse
+from typing import Any, Dict
 
 import numpy as np
 import requests
@@ -318,6 +319,32 @@ def download_from_cvat(self, max_retry_count=5, retry_wait_time=10):
             retry_count += 1
         return None
 
+    def add_directory(self, directory_id: str, filename: str, info: Dict[str, Any]) -> str:
+        """
+        Not implemented for this datastore
+
+        Abstract method for adding a directory to cvat
+        """
+        raise NotImplementedError("This datastore does not support adding directories")
+
+    def get_is_multichannel(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at multichannel (4D) data
+        """
+        logging.info("The function get_is_multichannel is not implemented for this datastore")
+        return False
+
+    def get_is_multi_file(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at directories containing multiple images per sample
+        """
+        logger.info("The function get_is_multi_file is not implemented for this datastore")
+        return False
+
 
 """
 def main():

monailabel/datastore/dicom.py

@@ -264,3 +264,29 @@ def _download_labeled_data(self):
     def datalist(self, full_path=True) -> List[Dict[str, Any]]:
         self._download_labeled_data()
         return super().datalist(full_path)
+
+    def add_directory(self, directory_id: str, filename: str, info: Dict[str, Any]) -> str:
+        """
+        Not implemented
+
+        Abstract method for adding a directory to DICOMWeb
+        """
+        raise NotImplementedError("This datastore does not support adding directories")
+
+    def get_is_multichannel(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at multichannel (4D) data
+        """
+        logging.info("The function get_is_multichannel is not implemented for this datastore")
+        return False
+
+    def get_is_multi_file(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at directories containing multiple images per sample
+        """
+        logger.info("The function get_is_multi_file is not implemented for this datastore")
+        return False

monailabel/datastore/dsa.py

@@ -270,6 +270,32 @@ def status(self) -> Dict[str, Any]:
     def json(self):
         return self.datalist()
 
+    def add_directory(self, directory_id: str, filename: str, info: Dict[str, Any]) -> str:
+        """
+        Not implemented for this datastore
+
+        Abstract method for adding a directory to dsa
+        """
+        raise NotImplementedError("This datastore does not support adding directories")
+
+    def get_is_multichannel(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at multichannel (4D) data
+        """
+        logging.info("The function get_is_multichannel is not implemented for this datastore")
+        return False
+
+    def get_is_multi_file(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at directories containing multiple images per sample
+        """
+        logger.info("The function get_is_multi_file is not implemented for this datastore")
+        return False
+
 
 """
 def main():

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: bool = False,
+        multi_file: bool = 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: bool = multichannel
+        self._multi_file: bool = multi_file
 
         logging.getLogger("filelock").setLevel(logging.ERROR)
 
@@ -256,6 +260,18 @@ 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:
+        """
+        Returns whether the dataset is multichannel or not
+        """
+        return self._multichannel
+
+    def get_is_multi_file(self) -> bool:
+        """
+        Returns whether the dataset is multichannel or not
+        """
+        return self._multi_file
+
     def get_image(self, image_id: str, params=None) -> Any:
         """
         Retrieve image object based on image id
@@ -431,6 +447,43 @@ def refresh(self):
         """
         self._reconcile_datastore()
 
+    def add_directory(self, directory_id: str, filename: str, info: Dict[str, Any]) -> str:
+        """
+        Add a directory to the datastore
+
+        :param directory_id: the directory id
+        :param filename: the filename
+        :param info: additional info
+
+        :return: directory id
+        """
+        id = os.path.basename(os.path.normpath(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!")
+            if os.path.isdir(filename):
+                if os.path.exists(dest):
+                    shutil.rmtree(dest)
+                shutil.copytree(filename, dest)
+            else:
+                shutil.copy2(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 +605,17 @@ 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:
+            ignored = {"labels", ".lock", os.path.basename(self._datastore_config_path).lower()}
+            for file in files:
+                abs_file = os.path.join(path, file)
+                if os.path.isdir(abs_file) and file.lower() not in ignored:
+                    filtered[os.path.basename(file)] = file
         return filtered
 
     def _reconcile_datastore(self):
@@ -585,24 +645,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(image=DataModel(info=file_info, ext=file_ext_str))
         return invalidate
 
     def _add_non_existing_labels(self, tag) -> int:

monailabel/datastore/xnat.py

@@ -386,6 +386,32 @@ def __upload_assessment(self, aiaa_model_name, image_id, file_path, type):
 
         self._request_put(url, data, type=type)
 
+    def add_directory(self, directory_id: str, filename: str, info: Dict[str, Any]) -> str:
+        """
+        Not implemented for this datastore
+
+        Abstract method for adding a directory to xnat
+        """
+        raise NotImplementedError("This datastore does not support adding directories")
+
+    def get_is_multichannel(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at multichannel (4D) data
+        """
+        logging.info("The function get_is_multichannel is not implemented for this datastore")
+        return False
+
+    def get_is_multi_file(self) -> bool:
+        """
+        Not implemented for this datastore
+
+        Returns whether the application's studies is directed at directories containing multiple images per sample
+        """
+        logger.info("The function get_is_multi_file is not implemented for this datastore")
+        return False
+
 
 """
 def main():

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,15 @@ 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}
+
+    if not os.path.isdir(image_file):
+        raise HTTPException(status_code=400, detail="Multi-file datastore requires a directory, not a file")
+
+    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: bool = strtobool(conf.get("multichannel", False))
+        self.multi_file: bool = strtobool(conf.get("multi_file", False))
+        self.input_channels = conf.get("input_channels", False)
         self.name = name
         self.description = description
         self.version = version
@@ -146,6 +148,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:
@@ -282,6 +286,9 @@ def infer(self, request, datastore=None):
             )
 
         request = copy.deepcopy(request)
+        request["multi_file"] = self.multi_file
+        request["multichannel"] = self.multichannel
+        request["input_channels"] = self.input_channels
         request["description"] = task.description
 
         image_id = request["image"]
@@ -292,7 +299,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']}")
@@ -431,6 +438,10 @@ def train(self, request):
             )
 
         request = copy.deepcopy(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
         result = task(request, self.datastore())
 
         # Run all scoring methods

monailabel/interfaces/datastore.py

@@ -201,6 +201,18 @@ def refresh(self) -> None:
         """
         pass
 
+    @abstractmethod
+    def add_directory(self, directory_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 directory_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 +291,17 @@ def json(self):
         Return json representation of datastore
         """
         pass
+
+    @abstractmethod
+    def get_is_multichannel(self) -> bool:
+        """
+        Returns whether the application's studies is directed at multichannel (4D) data
+        """
+        pass
+
+    @abstractmethod
+    def get_is_multi_file(self) -> bool:
+        """
+        Returns whether the application's studies is directed at directories containing multiple images per sample
+        """
+        pass