Skip to content

optionally use registered H&E images #37

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
NoraKrbr merged 1 commit into from
Nov 3, 2026
Merged

optionally use registered H&E images #37

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
120 changes: 93 additions & 27 deletions src/bio_image_datasets/pathocell2_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,37 +18,49 @@
7: 'Monocytes',
8: 'Granulocytes',
9: 'Endothelial Cells',
10: 'Smooth muscle Cells',
10: 'Smooth muscle Cells / Pericytes',
11: 'Fibroblasts',
12: 'Pericytes',
12: 'Pericytes (outdated)',
13: 'Dendritic Cells',
}


class PathoCell2Dataset(Dataset):
def __init__(self, local_path, version='v_0-0-1-beta'):
def __init__(self, local_path, version='v_0-0-1-beta', original_he_sample_names=None, registered=True):
"""
Initialize the PatchoCell2 dataset from TIFF files on disk.

The preliminary masks are stored under:
<local_path>/annotation/dataset_versions/<version>/masks
and <local_path>/annotation/dataset_versions/<version>/pheno

The H&E images are stored under:
The original H&E images are stored under:
<local_path>/he_image

The registered H&E images are stored under:
<local_path>/he_image_registered

The mIF images are stored under:
<local_path>/if_image

Args:
local_path (str): Root directory containing the dataset folders.
version (str): Dataset version folder name under annotation/dataset_versions.
original_he_sample_names (list[str] | None): Sample names that should use the
original (non-registered) H&E image instead of the registered one.
registered (bool): If True, use the registered H&E by default and only use the
original H&E for explicit sample names in ``original_he_sample_names``.
If False, use the original H&E for all samples.
"""
super().__init__(local_path)
self.local_path = os.path.expanduser(local_path)
self.version = version
self.registered = registered
self.original_he_sample_names = {self._normalize_sample_name(name) for name in (original_he_sample_names or ())}

self.he_registered_dir = os.path.join(self.local_path, 'he_image_registered')
self.he_dir = os.path.join(self.local_path, 'he_image')
self.he_dirs = self._resolve_he_dirs()
self.if_dir = os.path.join(self.local_path, 'if_image')
self.instance_dir = os.path.join(self.local_path, 'annotation', 'dataset_versions', self.version, 'masks')
self.semantic_dir = os.path.join(self.local_path, 'annotation', 'dataset_versions', self.version, 'pheno')
Expand All @@ -57,13 +69,36 @@ def __init__(self, local_path, version='v_0-0-1-beta'):
if not self.sample_names:
raise ValueError(
"No matching TIFF samples were found in the expected directories. "
f"Looked for files under: {self.he_dir}, {self.if_dir}, {self.instance_dir}, and {self.semantic_dir}."
f"Looked for files under: {self.he_dirs}, {self.if_dir}, {self.instance_dir}, and {self.semantic_dir}."
)

def _resolve_he_dirs(self):
"""Return the registered H&E directory first, then the original H&E directory as fallback."""
he_dirs = []
for directory_path in [self.he_registered_dir, self.he_dir]:
if os.path.isdir(directory_path):
he_dirs.append(directory_path)

if not he_dirs:
he_dirs.append(self.he_dir)
return he_dirs

def _discover_sample_names(self):
"""Collect sample names that exist in all required directories."""
directory_paths = [self.he_dir, self.if_dir, self.instance_dir, self.semantic_dir]
discovered = []
he_sample_stems = set()
for directory in self.he_dirs:
if not os.path.isdir(directory):
continue
for file_name in os.listdir(directory):
if not file_name.lower().endswith(('.tif', '.tiff', '.png')):
continue
stem = os.path.splitext(file_name)[0]
sample_stem = self._sample_stem_from_file_name(stem)
if sample_stem is not None:
he_sample_stems.add(sample_stem)

directory_paths = [self.if_dir, self.instance_dir, self.semantic_dir]
discovered = [he_sample_stems]
for directory in directory_paths:
if not os.path.isdir(directory):
raise FileNotFoundError(f"Required directory not found: {directory}")
Expand All @@ -82,30 +117,57 @@ def _discover_sample_names(self):
return common_names

@staticmethod
def _sample_stem_from_file_name(file_name):
"""Normalize sample names from different PatchoCell2 file naming conventions."""
def _normalize_sample_name(sample_name):
"""Normalize sample-name matching across case differences."""
return sample_name.lower()

@classmethod
def _sample_stem_from_file_name(cls, file_name):
"""Normalize PathoCell2 H&E file names to a canonical sample identifier."""
suffixes = [
'_cell_masks', '_nuclei_masks', '_cell_mask', '_nuclei_mask',
'_pheno', '_phenotype', '_mask', '_masks', '_he', '_if', '_simple_cell_types'
'_pheno', '_phenotype', '_mask', '_masks',
'_he_registered', '_he', '_if', '_simple_cell_types'
]
lowered_name = file_name.lower()
for suffix in suffixes:
if file_name.endswith(suffix):
return file_name[:-len(suffix)]
return file_name
if lowered_name.endswith(suffix):
return cls._normalize_sample_name(file_name[:-len(suffix)])
return cls._normalize_sample_name(file_name)

def _get_sample_path(self, directory, sample_name):
"""Return the first matching file for a sample in a directory."""
candidates = []
for file_name in os.listdir(directory):
if file_name.lower().endswith(('.tif', '.tiff', '.png')):
stem = os.path.splitext(file_name)[0]
if self._sample_stem_from_file_name(stem) == sample_name:
candidates.append(os.path.join(directory, file_name))

if not candidates:
raise FileNotFoundError(f"No file for sample '{sample_name}' found in {directory}")

return sorted(candidates)[0]
"""Return the first matching file for a sample in the given directory priority order."""
if isinstance(directory, (list, tuple)):
directories = list(directory)
else:
directories = [directory]

for current_directory in directories:
for file_name in sorted(os.listdir(current_directory)):
if file_name.lower().endswith(('.tif', '.tiff', '.png')):
stem = os.path.splitext(file_name)[0]
if self._sample_stem_from_file_name(stem) == sample_name:
return os.path.join(current_directory, file_name)

search_dirs = ', '.join(str(path) for path in directories)
raise FileNotFoundError(f"No file for sample '{sample_name}' found in {search_dirs}")

def _get_he_path(self, sample_name):
"""Select the H&E image according to the registration mode and sample-specific overrides."""
normalized_name = self._normalize_sample_name(sample_name)

if self.registered:
preferred_dirs = [self.he_registered_dir, self.he_dir]
if normalized_name in self.original_he_sample_names and os.path.isdir(self.he_dir):
preferred_dirs = [self.he_dir, self.he_registered_dir]
else:
preferred_dirs = [self.he_dir, self.he_registered_dir]

preferred_dirs = [directory for directory in preferred_dirs if os.path.isdir(directory)]
if not preferred_dirs:
preferred_dirs = [self.he_dir]

return self._get_sample_path(preferred_dirs, normalized_name)

def _get_instance_mask_paths(self, sample_name):
"""Return both cell and nuclei instance masks for a sample, if present."""
Expand All @@ -127,7 +189,11 @@ def _get_instance_mask_paths(self, sample_name):
def _load_image(self, directory, idx):
"""Load an image from TIFF and return it in channel-first layout."""
sample_name = self.sample_names[idx]
image = np.asarray(io.imread(self._get_sample_path(directory, sample_name)))
if directory == self.he_dirs:
path = self._get_he_path(sample_name)
else:
path = self._get_sample_path(directory, sample_name)
image = np.asarray(io.imread(path))

if image.ndim == 2:
return image[np.newaxis, :, :]
Expand Down Expand Up @@ -166,7 +232,7 @@ def get_he(self, idx):
"""Load the H&E image at the given index."""
if idx >= len(self):
raise IndexError("Index out of bounds.")
return self._load_image(self.he_dir, idx)
return self._load_image(self.he_dirs, idx)

def get_if(self, idx):
"""Load the mIF image at the given index."""
Expand Down
136 changes: 136 additions & 0 deletions tests/test_pathocell2_dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,59 @@ def prepare_pathocell2_samples(output_dir, num_samples=4, size=(32, 32)):
return sample_names


def prepare_pathocell2_registered_he_samples(output_dir, num_samples=4, size=(32, 32)):
"""Create mock TIFF files where the registered H&E images are the preferred default."""
registered_he_dir = os.path.join(output_dir, "he_image_registered")
if_dir = os.path.join(output_dir, "if_image")
masks_dir = os.path.join(output_dir, "annotation", "dataset_versions", "v_0-0-1-beta", "masks")
pheno_dir = os.path.join(output_dir, "annotation", "dataset_versions", "v_0-0-1-beta", "pheno")

for directory in [registered_he_dir, if_dir, masks_dir, pheno_dir]:
os.makedirs(directory, exist_ok=True)

sample_names = [f"sample_{index}" for index in range(num_samples)]
for sample_name in sample_names:
registered_he_image = np.full((size[0], size[1], 3), 7, dtype=np.uint8)
if_image = np.random.randint(0, 256, size=(size[0], size[1], 4), dtype=np.uint8)
instance_mask = np.random.randint(0, 1000, size=size, dtype=np.uint16)
semantic_mask = np.random.randint(0, 14, size=size, dtype=np.uint8)

io.imsave(os.path.join(registered_he_dir, f"{sample_name}_he_registered.tif"), registered_he_image)
io.imsave(os.path.join(if_dir, f"{sample_name}_if.tif"), if_image)
io.imsave(os.path.join(masks_dir, f"{sample_name}_cell_masks.tif"), instance_mask)
io.imsave(os.path.join(pheno_dir, f"{sample_name}_pheno.tif"), semantic_mask)

return sample_names


def prepare_pathocell2_non_registered_he_override_samples(output_dir, num_samples=4, size=(32, 32)):
"""Create mock TIFF files where one sample explicitly opts back to the non-registered H&E."""
registered_he_dir = os.path.join(output_dir, "he_image_registered")
legacy_he_dir = os.path.join(output_dir, "he_image")
if_dir = os.path.join(output_dir, "if_image")
masks_dir = os.path.join(output_dir, "annotation", "dataset_versions", "v_0-0-1-beta", "masks")
pheno_dir = os.path.join(output_dir, "annotation", "dataset_versions", "v_0-0-1-beta", "pheno")

for directory in [registered_he_dir, legacy_he_dir, if_dir, masks_dir, pheno_dir]:
os.makedirs(directory, exist_ok=True)

sample_names = [f"sample_{index}" for index in range(num_samples)]
for sample_name in sample_names:
registered_he_image = np.full((size[0], size[1], 3), 7, dtype=np.uint8)
legacy_he_image = np.full((size[0], size[1], 3), 9, dtype=np.uint8)
if_image = np.random.randint(0, 256, size=(size[0], size[1], 4), dtype=np.uint8)
instance_mask = np.random.randint(0, 1000, size=size, dtype=np.uint16)
semantic_mask = np.random.randint(0, 14, size=size, dtype=np.uint8)

io.imsave(os.path.join(registered_he_dir, f"{sample_name}_he_registered.tif"), registered_he_image)
io.imsave(os.path.join(legacy_he_dir, f"{sample_name}_he_non_registered.tif"), legacy_he_image)
io.imsave(os.path.join(if_dir, f"{sample_name}_if.tif"), if_image)
io.imsave(os.path.join(masks_dir, f"{sample_name}_cell_masks.tif"), instance_mask)
io.imsave(os.path.join(pheno_dir, f"{sample_name}_pheno.tif"), semantic_mask)

return sample_names


def test_len():
with tempfile.TemporaryDirectory() as tmp_dir:
prepare_pathocell2_samples(tmp_dir, num_samples=4)
Expand Down Expand Up @@ -70,6 +123,89 @@ def test_get_he():
assert he_image.shape == (3, 32, 32)


def test_get_he_prefers_registered_images_by_default():
with tempfile.TemporaryDirectory() as tmp_dir:
prepare_pathocell2_registered_he_samples(tmp_dir, num_samples=4)
dataset = PathoCell2Dataset(local_path=tmp_dir)
he_image = dataset.get_he(0)
assert he_image.shape == (3, 32, 32)
assert he_image.dtype == np.uint8
assert np.all(he_image == 7)


def test_get_he_uses_original_he_for_explicit_sample_names():
with tempfile.TemporaryDirectory() as tmp_dir:
sample_names = prepare_pathocell2_non_registered_he_override_samples(tmp_dir, num_samples=4)
original_he_dir = os.path.join(tmp_dir, 'he_image')
os.makedirs(original_he_dir, exist_ok=True)
for sample_name in sample_names[:3]:
original_he_image = np.full((32, 32, 3), 9, dtype=np.uint8)
io.imsave(os.path.join(original_he_dir, f'{sample_name}_he.tif'), original_he_image)

original_samples = sample_names[:3]
dataset = PathoCell2Dataset(local_path=tmp_dir, original_he_sample_names=original_samples)

for idx, sample_name in enumerate(original_samples):
he_image = dataset.get_he(dataset.get_sample_names().index(sample_name))
assert np.all(he_image == 9)

default_he_image = dataset.get_he(3)
assert np.all(default_he_image == 7)


def test_get_he_uses_requested_original_sample_names_for_reg_files():
with tempfile.TemporaryDirectory() as tmp_dir:
registered_he_dir = os.path.join(tmp_dir, 'he_image_registered')
original_he_dir = os.path.join(tmp_dir, 'he_image')
if_dir = os.path.join(tmp_dir, 'if_image')
masks_dir = os.path.join(tmp_dir, 'annotation', 'dataset_versions', 'v_0-0-1-beta', 'masks')
pheno_dir = os.path.join(tmp_dir, 'annotation', 'dataset_versions', 'v_0-0-1-beta', 'pheno')
for directory in [registered_he_dir, original_he_dir, if_dir, masks_dir, pheno_dir]:
os.makedirs(directory, exist_ok=True)

registered_he = np.full((32, 32, 3), 7, dtype=np.uint8)
original_he = np.full((32, 32, 3), 9, dtype=np.uint8)
if_image = np.random.randint(0, 256, size=(32, 32, 4), dtype=np.uint8)
instance_mask = np.random.randint(0, 1000, size=(32, 32), dtype=np.uint16)
semantic_mask = np.random.randint(0, 14, size=(32, 32), dtype=np.uint8)

io.imsave(os.path.join(registered_he_dir, 'reg001_a_he_registered.tiff'), registered_he)
io.imsave(os.path.join(original_he_dir, 'reg001_A_he.tiff'), original_he)
io.imsave(os.path.join(if_dir, 'reg001_a_if.tiff'), if_image)
io.imsave(os.path.join(masks_dir, 'reg001_a_cell_masks.tiff'), instance_mask)
io.imsave(os.path.join(pheno_dir, 'reg001_a_pheno.tiff'), semantic_mask)

dataset = PathoCell2Dataset(local_path=tmp_dir, original_he_sample_names=['reg001_A'])
assert np.all(dataset.get_he(0) == 9)
assert dataset.get_sample_name(0) == 'reg001_a'


def test_get_he_uses_original_for_all_samples_when_registered_false():
with tempfile.TemporaryDirectory() as tmp_dir:
registered_he_dir = os.path.join(tmp_dir, 'he_image_registered')
original_he_dir = os.path.join(tmp_dir, 'he_image')
if_dir = os.path.join(tmp_dir, 'if_image')
masks_dir = os.path.join(tmp_dir, 'annotation', 'dataset_versions', 'v_0-0-1-beta', 'masks')
pheno_dir = os.path.join(tmp_dir, 'annotation', 'dataset_versions', 'v_0-0-1-beta', 'pheno')
for directory in [registered_he_dir, original_he_dir, if_dir, masks_dir, pheno_dir]:
os.makedirs(directory, exist_ok=True)

registered_he = np.full((32, 32, 3), 7, dtype=np.uint8)
original_he = np.full((32, 32, 3), 9, dtype=np.uint8)
if_image = np.random.randint(0, 256, size=(32, 32, 4), dtype=np.uint8)
instance_mask = np.random.randint(0, 1000, size=(32, 32), dtype=np.uint16)
semantic_mask = np.random.randint(0, 14, size=(32, 32), dtype=np.uint8)

io.imsave(os.path.join(registered_he_dir, 'reg001_a_he_registered.tiff'), registered_he)
io.imsave(os.path.join(original_he_dir, 'reg001_A_he.tiff'), original_he)
io.imsave(os.path.join(if_dir, 'reg001_a_if.tiff'), if_image)
io.imsave(os.path.join(masks_dir, 'reg001_a_cell_masks.tiff'), instance_mask)
io.imsave(os.path.join(pheno_dir, 'reg001_a_pheno.tiff'), semantic_mask)

dataset = PathoCell2Dataset(local_path=tmp_dir, registered=False)
assert np.all(dataset.get_he(0) == 9)


def test_get_if():
with tempfile.TemporaryDirectory() as tmp_dir:
prepare_pathocell2_samples(tmp_dir, num_samples=4)
Expand Down
Loading