@@ -102,9 +102,11 @@ def __init__(
102102 images_dir : str = "." ,
103103 labels_dir : str = "labels" ,
104104 datastore_config : str = "datastore_v2.json" ,
105- extensions = ("*.nii.gz" , "*.nii" ),
105+ extensions = ("*.nii.gz" , "*.nii" , "*.nrrd" ),
106106 auto_reload = False ,
107107 read_only = False ,
108+ multichannel : bool = False ,
109+ multi_file : bool = False ,
108110 ):
109111 """
110112 Creates a `LocalDataset` object
@@ -124,6 +126,8 @@ def __init__(
124126 self ._ignore_event_config = False
125127 self ._config_ts = 0
126128 self ._auto_reload = auto_reload
129+ self ._multichannel : bool = multichannel
130+ self ._multi_file : bool = multi_file
127131
128132 logging .getLogger ("filelock" ).setLevel (logging .ERROR )
129133
@@ -256,6 +260,12 @@ def datalist(self, full_path=True) -> List[Dict[str, Any]]:
256260 ds = json .loads (json .dumps (ds ).replace (f"{ self ._datastore_path .rstrip (os .pathsep )} { os .pathsep } " , "" ))
257261 return ds
258262
263+ def get_is_multichannel (self ) -> bool :
264+ return self ._multichannel
265+
266+ def get_is_multi_file (self ) -> bool :
267+ return self ._multi_file
268+
259269 def get_image (self , image_id : str , params = None ) -> Any :
260270 """
261271 Retrieve image object based on image id
@@ -431,6 +441,29 @@ def refresh(self):
431441 """
432442 self ._reconcile_datastore ()
433443
444+ def add_directory (self , directory_id : str , filename : str , info : Dict [str , Any ]) -> str :
445+ id = os .path .basename (filename )
446+ if not directory_id :
447+ directory_id = id
448+
449+ logger .info (f"Adding Image: { directory_id } => { filename } " )
450+ name = directory_id
451+ dest = os .path .realpath (os .path .join (self ._datastore .image_path (), name ))
452+
453+ with FileLock (self ._lock_file ):
454+ logger .debug ("Acquired the lock!" )
455+ shutil .copy (filename , dest )
456+
457+ info = info if info else {}
458+ info ["ts" ] = int (time .time ())
459+ info ["name" ] = name
460+
461+ # images = get_directory_contents(filename)
462+ self ._datastore .objects [directory_id ] = ImageLabelModel (image = DataModel (info = info , ext = "" ))
463+ self ._update_datastore_file (lock = False )
464+ logger .debug ("Released the lock!" )
465+ return directory_id
466+
434467 def add_image (self , image_id : str , image_filename : str , image_info : Dict [str , Any ]) -> str :
435468 id , image_ext = self ._to_id (os .path .basename (image_filename ))
436469 if not image_id :
@@ -552,10 +585,15 @@ def _list_files(self, path, patterns):
552585 files = os .listdir (path )
553586
554587 filtered = dict ()
555- for pattern in patterns :
556- matching = fnmatch .filter (files , pattern )
557- for file in matching :
558- filtered [os .path .basename (file )] = file
588+ if not self ._multi_file :
589+ for pattern in patterns :
590+ matching = fnmatch .filter (files , pattern )
591+ for file in matching :
592+ filtered [os .path .basename (file )] = file
593+ else :
594+ for file in files :
595+ if file .lower () not in ["labels" , ".lock" , "datastore_v2.json" ]:
596+ filtered [os .path .basename (file )] = file
559597 return filtered
560598
561599 def _reconcile_datastore (self ):
@@ -585,23 +623,26 @@ def _add_non_existing_images(self) -> int:
585623 invalidate = 0
586624 self ._init_from_datastore_file ()
587625
588- local_images = self ._list_files (self ._datastore .image_path (), self ._extensions )
626+ local_files = self ._list_files (self ._datastore .image_path (), self ._extensions )
589627
590- image_ids = list (self ._datastore .objects .keys ())
591- for image_file in local_images :
592- image_id , image_ext = self ._to_id (image_file )
593- if image_id not in image_ids :
594- logger .info (f"Adding New Image: { image_id } => { image_file } " )
628+ ids = list (self ._datastore .objects .keys ())
629+ for file in local_files :
630+ if self ._multi_file :
631+ # Directories have no extension — use the name as-is
632+ file_id = file
633+ file_ext_str = ""
634+ else :
635+ file_id , file_ext_str = self ._to_id (file )
595636
596- name = self ._filename (image_id , image_ext )
597- image_info = {
637+ if file_id not in ids :
638+ logger .info (f"Adding New Image: { file_id } => { file } " )
639+ name = self ._filename (file_id , file_ext_str )
640+ file_info = {
598641 "ts" : int (time .time ()),
599- # "checksum": file_checksum(os.path.join(self._datastore.image_path(), name)),
600642 "name" : name ,
601643 }
602-
603644 invalidate += 1
604- self ._datastore .objects [image_id ] = ImageLabelModel (image = DataModel (info = image_info , ext = image_ext ))
645+ self ._datastore .objects [file_id ] = ImageLabelModel (name = DataModel (info = file_info , ext = file_ext_str ))
605646
606647 return invalidate
607648
0 commit comments