add support for tensorstore

Author: donglaiw

connectomics/config/defaults.py

@@ -474,6 +474,7 @@
 _C.INFERENCE.INPUT_SIZE = None
 _C.INFERENCE.OUTPUT_SIZE = None
 
+_C.INFERENCE.TENSORSTORE_PATH = None
 _C.INFERENCE.INPUT_PATH = None
 _C.INFERENCE.IMAGE_NAME = None
 _C.INFERENCE.OUTPUT_PATH = ""

connectomics/data/dataset/build.py

@@ -208,6 +208,7 @@ def _validate_shape(cfg, image, mask, i):
 
     for i in range(num_vols):
         if volume is not None:
+
             volume[i] = read_fn(img_name[i], drop_channel=cfg.DATASET.DROP_CHANNEL)
             print(f"volume shape (original): {volume[i].shape}")
             if cfg.DATASET.NORMALIZE_RANGE:
@@ -255,7 +256,9 @@ def get_dataset(cfg,
                 dataset_class=VolumeDataset,
                 dataset_options={},
                 dir_name_init: Optional[list] = None,
-                img_name_init: Optional[list] = None):
+                img_name_init: Optional[list] = None,
+                tensorstore_data = None,
+                tensorstore_coord: Optional[list] = None):
     r"""Prepare dataset for training and inference.
     """
     assert mode in ['train', 'val', 'test']
@@ -337,8 +340,14 @@ def _make_json_path(path, name):
                               **shared_kwargs)
 
     else:  # build VolumeDataset or VolumeDatasetMultiSeg
-        volume, label, valid_mask = _get_input(
-            cfg, mode, rank, dir_name_init, img_name_init, min_size=sample_volume_size)
+        if tensorstore_data is None: 
+            volume, label, valid_mask = _get_input(
+                cfg, mode, rank, dir_name_init, img_name_init, min_size=sample_volume_size)
+        else:
+            volume = [tensorstore_data[coord[0]:coord[1],coord[2]:coord[3],coord[4]:coord[5]].read().result().transpose() \
+                        for coord in tensorstore_coord]
+            label = None
+            valid_mask = None
 
         if cfg.MODEL.TARGET_OPT_MULTISEG_SPLIT is not None:
             shared_kwargs['multiseg_split'] = cfg.MODEL.TARGET_OPT_MULTISEG_SPLIT

connectomics/data/utils/data_io.py

@@ -11,6 +11,7 @@
 import glob
 import numpy as np
 import imageio
+import pickle
 from scipy.ndimage import zoom
 
 
@@ -110,6 +111,24 @@ def readimgs(filename):
 
     return data
 
+def read_pkl(filename):
+    """
+    The function `read_pkl` reads a pickle file and returns a list of the objects stored in the file.
+
+    :param filename: The filename parameter is a string that represents the name of the file you want to
+    read. It should include the file extension, such as ".pkl" for a pickle file
+    :return: a list of objects that were read from the pickle file.
+    """
+    data = []
+    with open(filename, "rb") as fid:
+        while True:
+            try:
+                data.append(pickle.load(fid))
+            except EOFError:
+                break
+    if len(data) == 1:
+        return data[0]
+    return data
 
 def writeh5(filename, dtarray, dataset='main'):
     fid = h5py.File(filename, 'w')

connectomics/engine/trainer.py

@@ -5,6 +5,7 @@
 import os
 import time
 import math
+import pickle
 import GPUtil
 import numpy as np
 from yacs.config import CfgNode
@@ -19,7 +20,7 @@
 from ..data.augmentation import build_train_augmentor, TestAugmentor
 from ..data.dataset import build_dataloader, get_dataset
 from ..data.dataset.build import _get_file_list
-from ..data.utils import build_blending_matrix, writeh5
+from ..data.utils import build_blending_matrix, writeh5, read_pkl
 from ..data.utils import get_padsize, array_unpad
 
 
@@ -272,32 +273,53 @@ def test(self):
             writeh5(save_path, result, ['vol%d' % (x) for x in range(len(result))])
             print('Prediction saved as: ', save_path)
 
-    def test_singly(self):
-        dir_name = _get_file_list(self.cfg.DATASET.INPUT_PATH)
-        assert len(dir_name) == 1 # avoid ambiguity when DO_SINGLY is True
-        img_name = _get_file_list(self.cfg.DATASET.IMAGE_NAME, prefix=dir_name[0])
-        num_file = len(img_name)
-
-        if os.path.isfile(self.cfg.INFERENCE.OUTPUT_NAME):
-            output_name = _get_file_list(self.cfg.DATASET.OUTPUT_NAME, prefix=self.output_dir)
+    def test_singly(self):        
+        dir_name = None
+        if self.cfg.INFERENCE.TENSORSTORE_PATH is None:
+            dir_name = _get_file_list(self.cfg.DATASET.INPUT_PATH)
+            assert len(dir_name) == 1 # avoid ambiguity when DO_SINGLY is True
+            img_name = _get_file_list(self.cfg.DATASET.IMAGE_NAME, prefix=dir_name[0])
         else:
-            # same filename but different location
-            if self.output_dir != dir_name[0]:
-                output_name = _get_file_list(self.cfg.DATASET.IMAGE_NAME, prefix=self.output_dir)
+            import tensorstore as ts
+            context = ts.Context({'cache_pool': {'total_bytes_limit': 1000000000}})
+            ts_dict = read_pkl(self.cfg.INFERENCE.TENSORSTORE_PATH)
+            ts_data = ts.open(ts_dict, read=True, context=context).result()[ts.d['channel'][0]]
+            # chunk coordinate
+            img_name = np.loadtxt(self.cfg.DATASET.IMAGE_NAME).astype(int)
+                            
+        num_file = len(img_name)
+        
+        if os.path.isfile(os.path.join(self.output_dir, self.cfg.INFERENCE.OUTPUT_NAME)):
+            # load output names 
+            output_name = _get_file_list(self.cfg.INFERENCE.OUTPUT_NAME, prefix=self.output_dir)
+        else:            
+            if dir_name is None or self.output_dir != dir_name[0]:
+                # same filenames but different location
+                if '{' in self.cfg.INFERENCE.OUTPUT_NAME:
+                    # template function
+                    output_name = [None] * num_file
+                    for i in range(num_file):
+                        arr = img_name[i]
+                        output_name[i] = os.path.join(self.output_dir, eval(self.cfg.INFERENCE.OUTPUT_NAME)+'.h5')
+                else: 
+                    output_name = _get_file_list(self.cfg.DATASET.IMAGE_NAME, prefix=self.output_dir)
             else:
+                # same file location
                 output_name = [x+'_result.h5' for x in img_name]
 
-        # save input image names for future reference
-        fw = open(os.path.join(self.output_dir, "images.txt"), "w")
-        fw.write('\n'.join(img_name))
-        fw.close()
-
         for i in range(self.cfg.INFERENCE.DO_SINGLY_START_INDEX, num_file, self.cfg.INFERENCE.DO_SINGLY_STEP):
             self.test_filename = output_name[i]
             if not os.path.exists(self.test_filename):
-                dataset = get_dataset(
-                    self.cfg, self.augmentor, self.mode, self.rank,
-                    dir_name_init=dir_name, img_name_init=[img_name[i]])
+                if self.cfg.INFERENCE.TENSORSTORE_PATH is None:
+                    # directly load from dir_name_init and img_name_init
+                    dataset = get_dataset(
+                        self.cfg, self.augmentor, self.mode, self.rank,
+                        dir_name_init=dir_name, img_name_init=[img_name[i]])
+                else:
+                    # preload from tensorstore
+                    dataset = get_dataset(
+                        self.cfg, self.augmentor, self.mode, self.rank, 
+                        tensorstore_data=ts_data, tensorstore_coord=[img_name[i]])
                 self.dataloader = build_dataloader(
                     self.cfg, self.augmentor, self.mode, dataset, self.rank)
                 self.dataloader = iter(self.dataloader)