Skip to content

Commit

Permalink
Merge pull request #95 from sameeul/bioformat_fix2
Browse files Browse the repository at this point in the history 
Make BioReader pickleable
  • Loading branch information
@sameeul
sameeul authored Aug 12, 2024
2 parents b7bb145 + db6a066 commit 6c99304
Show file tree
Hide file tree
Showing 4 changed files with 124 additions and 55 deletions.
16 changes: 12 additions & 4 deletions src/bfio/backends.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -122,18 +122,26 @@ def __init__(self, frontend):
def __getstate__(self) -> Dict:
state_dict = {n: getattr(self, n) for n in self._STATE_DICT}
state_dict.update({"file_path": self.frontend._file_path})
state_dict.update({"level": self.frontend.level})

return state_dict

def __setstate__(self, state) -> None:
for k, v in state.items():
if k == "file_path":
self._lock = threading.Lock()
self._rdr = tifffile.TiffFile(v)
self._rdr.filehandle.close()
if k == "file_path" or k == "level":
pass
else:
setattr(self, k, v)

self._lock = threading.Lock()
self._rdr = tifffile.TiffFile(state["file_path"])
self._rdr_pages = self._rdr.pages
if state["level"] is not None:
if len(self._rdr.series) != 0:
series = self._rdr.series[0]
self._rdr_pages = series.levels[state["level"]]
self._rdr.filehandle.close()

def read_metadata(self):
self.logger.debug("read_metadata(): Reading metadata...")

Expand Down
26 changes: 20 additions & 6 deletions src/bfio/bfio.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,16 +43,15 @@ class BioReader(BioBase):

logger = logging.getLogger("bfio.bfio.BioReader")
_STATE_DICT = [
"level",
"_metadata",
"_DIMS",
"_file_path",
"max_workers",
"_max_workers",
"_backend_name",
"clean_metadata",
"_read_only",
"_backend",
"level",
"append",
]

def __init__(
Expand Down Expand Up @@ -234,18 +233,33 @@ def set_backend(self, backend: typing.Optional[str] = None) -> None:
self._backend_name = backend.lower()

def __getstate__(self) -> typing.Dict:
state_dict = {n: getattr(self, n) for n in self._STATE_DICT}
if self._backend_name == "bioformats":
state_dict = {}
for n in self._STATE_DICT:
if n == "_backend":
state_dict[n] = "JavaReaderDummy"
else:
state_dict[n] = getattr(self, n)
else:
state_dict = {n: getattr(self, n) for n in self._STATE_DICT}

return state_dict

def __setstate__(self, state) -> None:
assert all(n in self._STATE_DICT for n in state.keys())
assert all(n in state.keys() for n in self._STATE_DICT)

bioformats_backend = False
for k, v in state.items():
setattr(self, k, v)
if k == "_backend" and v == "JavaReaderDummy":
bioformats_backend = True
else:
setattr(self, k, v)

self._backend.frontend = self
if bioformats_backend:
self._backend = backends.JavaReader(self)
else:
self._backend.frontend = self

def __getitem__(self, keys: typing.Union[tuple, slice]) -> numpy.ndarray:
"""Image loading using numpy-like indexing.
Expand Down
28 changes: 19 additions & 9 deletions src/bfio/ts_backends.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,19 @@ class TensorstoreReader(bfio.base_classes.TSAbstractReader):

_rdr = None
_offsets_bytes = None
_STATE_DICT = ["_metadata", "frontend"]
_STATE_DICT = [
"_metadata",
"frontend",
"X",
"Y",
"Z",
"C",
"T",
"data_type",
"_file_path",
"_file_type",
"_axes_list",
]

def __init__(self, frontend):
super().__init__(frontend)
Expand All @@ -33,18 +45,19 @@ def __init__(self, frontend):
if extension.endswith(".ome.tif") or extension.endswith(".ome.tiff"):
# # check if it satisfies all the condition for python backend
self._file_type = FileType.OmeTiff
self._rdr = TSReader(str(self.frontend._file_path), FileType.OmeTiff, "")
self._file_path = str(self.frontend._file_path.resolve())
self._axes_list = ""
elif extension.endswith(".zarr"):
# if path exists, make sure it is a directory
if not Path.is_dir(self.frontend._file_path):
raise ValueError(
"this filetype is not supported by tensorstore backend"
)
else:
zarr_path, axes_list = self.get_zarr_array_info()
self._file_path, self._axes_list = self.get_zarr_array_info()
self._file_type = FileType.OmeZarr
self._rdr = TSReader(zarr_path, FileType.OmeZarr, axes_list)

self._rdr = TSReader(self._file_path, self._file_type, self._axes_list)
self.X = self._rdr._X
self.Y = self._rdr._Y
self.Z = self._rdr._Z
Expand Down Expand Up @@ -141,16 +154,13 @@ def get_zarr_array_info(self):

def __getstate__(self) -> Dict:
state_dict = {n: getattr(self, n) for n in self._STATE_DICT}
state_dict.update({"file_path": self.frontend._file_path})

return state_dict

def __setstate__(self, state) -> None:
for k, v in state.items():
if k == "file_path":
pass
else:
setattr(self, k, v)
setattr(self, k, v)
self._rdr = TSReader(self._file_path, self._file_type, self._axes_list)

def read_metadata(self):

Expand Down
109 changes: 73 additions & 36 deletions tests/test_read.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import requests, pathlib, shutil, logging, sys
import bfio
import numpy as np
import pickle
import random
import zarr
from ome_zarr.utils import download as zarr_download
Expand Down Expand Up @@ -64,6 +65,19 @@ def setUpModule():
for z in range(br.Z):
zf[t, c, z, :, :] = br[:, :, z, c, t]

# create a 2D numpy array filled with random integer form 0-255
img_height = 8000
img_width = 7500
source_data = np.random.randint(0, 256, (img_height, img_width), dtype=np.uint16)
np.save(TEST_DIR.joinpath("random_image.npy"), source_data)
with bfio.BioWriter(
str(TEST_DIR.joinpath("random_image.ome.tiff")),
X=img_width,
Y=img_height,
dtype=np.uint16,
) as bw:
bw[0:img_height, 0:img_width, 0, 0, 0] = source_data


def tearDownModule():
"""Remove test images"""
Expand Down Expand Up @@ -141,24 +155,10 @@ def test_read_tif_strip_python(self):
I = br[:]

def test_read_unaligned_tile_boundary_python(self):
# create a 2D numpy array filled with random integer form 0-255
img_height = 8000
img_width = 7500
source_data = np.random.randint(
0, 256, (img_height, img_width), dtype=np.uint16
)
with bfio.BioWriter(
str(TEST_DIR.joinpath("test_output.ome.tiff")),
X=img_width,
Y=img_height,
dtype=np.uint16,
) as bw:
bw[0:img_height, 0:img_width, 0, 0, 0] = source_data

source_data = np.load(TEST_DIR.joinpath("random_image.npy"))
x_max = source_data.shape[0]
y_max = source_data.shape[1]

with bfio.BioReader(str(TEST_DIR.joinpath("test_output.ome.tiff"))) as test_br:
with bfio.BioReader(str(TEST_DIR.joinpath("random_image.ome.tiff"))) as br:
for i in range(100):
x_start = random.randint(0, x_max)
y_start = random.randint(0, y_max)
Expand All @@ -167,39 +167,25 @@ def test_read_unaligned_tile_boundary_python(self):

x_end = x_start + x_step
y_end = y_start + y_step

if x_end > x_max:
x_end = x_max

if y_end > y_max:
y_end = y_max

test_data = test_br[x_start:x_end, y_start:y_end, ...]
test_data = br[x_start:x_end, y_start:y_end, ...]
assert (
np.sum(source_data[x_start:x_end, y_start:y_end] - test_data) == 0
)

def test_read_unaligned_tile_boundary_tensorstore(self):
# create a 2D numpy array filled with random integer form 0-255
img_height = 8000
img_width = 7500
source_data = np.random.randint(
0, 256, (img_height, img_width), dtype=np.uint16
)
with bfio.BioWriter(
str(TEST_DIR.joinpath("test_output.ome.tiff")),
X=img_width,
Y=img_height,
dtype=np.uint16,
) as bw:
bw[0:img_height, 0:img_width, 0, 0, 0] = source_data

source_data = np.load(TEST_DIR.joinpath("random_image.npy"))
x_max = source_data.shape[0]
y_max = source_data.shape[1]

with bfio.BioReader(
str(TEST_DIR.joinpath("test_output.ome.tiff")), backend="tensorstore"
) as test_br:
str(TEST_DIR.joinpath("random_image.ome.tiff")), backend="tensorstore"
) as br:
for i in range(100):
x_start = random.randint(0, x_max)
y_start = random.randint(0, y_max)
Expand All @@ -208,14 +194,13 @@ def test_read_unaligned_tile_boundary_tensorstore(self):

x_end = x_start + x_step
y_end = y_start + y_step

if x_end > x_max:
x_end = x_max

if y_end > y_max:
y_end = y_max

test_data = test_br[x_start:x_end, y_start:y_end, ...]
test_data = br[x_start:x_end, y_start:y_end, ...]
assert (
np.sum(source_data[x_start:x_end, y_start:y_end] - test_data) == 0
)
Expand Down Expand Up @@ -305,6 +290,21 @@ def test_sub_resolution_read(self):
get_dims(br)
self.assertEqual(br.shape, (1167, 404, 1, 3))

def test_bioformats_backend_pickle(self):

br = bfio.BioReader(str(TEST_DIR.joinpath("img_r001_c001.ome.tif")))
img_height = br.Y
img_width = br.X
data_sum = br[:].sum()
pickled_rdr = pickle.dumps(br)
br.close()
unpickled_rdr = pickle.loads(pickled_rdr)

assert unpickled_rdr.X == img_width
assert unpickled_rdr.Y == img_height
assert unpickled_rdr[:].sum() == data_sum
unpickled_rdr.close()


class TestZarrReader(unittest.TestCase):
def test_get_dims(self):
Expand Down Expand Up @@ -420,3 +420,40 @@ def test_set_metadata(self):
logger.info(br.cnames)
logger.info(br.ps_x)
self.assertEqual(br.cnames[0], cname[0])


class TestBioReaderPickle(unittest.TestCase):

def test_python_backend_pickle(self):

br = bfio.BioReader(
str(TEST_DIR.joinpath("random_image.ome.tiff")), backend="python"
)
img_height = br.Y
img_width = br.X
data_sum = br[:].sum()
pickled_rdr = pickle.dumps(br)
br.close()
unpickled_rdr = pickle.loads(pickled_rdr)

assert unpickled_rdr.X == img_width
assert unpickled_rdr.Y == img_height
assert unpickled_rdr[:].sum() == data_sum
unpickled_rdr.close()

def test_tensorstore_backend_pickle(self):

br = bfio.BioReader(
str(TEST_DIR.joinpath("random_image.ome.tiff")), backend="tensorstore"
)
img_height = br.Y
img_width = br.X
data_sum = br[:].sum()
pickled_rdr = pickle.dumps(br)
br.close()
unpickled_rdr = pickle.loads(pickled_rdr)

assert unpickled_rdr.X == img_width
assert unpickled_rdr.Y == img_height
assert unpickled_rdr[:].sum() == data_sum
unpickled_rdr.close()

0 comments on commit 6c99304

Please sign in to comment.