Spaces:
Running
Running
File size: 2,865 Bytes
a4b3c40 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import numpy as np
import copy
from monai.transforms import RandScaleIntensity, Compose
from monai.transforms.compose import MapTransform
from skimage.segmentation import find_boundaries
__all__ = ["BoundaryExclusion", "IntensityDiversification"]
class BoundaryExclusion(MapTransform):
"""Map the cell boundary pixel labels to the background class (0)."""
def __init__(self, keys=["label"], allow_missing_keys=False):
super(BoundaryExclusion, self).__init__(keys, allow_missing_keys)
def __call__(self, data):
# Find and Exclude Boundary
label_original = data["label"]
label = copy.deepcopy(label_original)
boundary = find_boundaries(label, connectivity=1, mode="thick")
label[boundary] = 0
# Do not exclude if the cell is too small (< 14x14).
new_label = copy.deepcopy(label_original)
new_label[label == 0] = 0
cell_idx, cell_counts = np.unique(label_original, return_counts=True)
for k in range(len(cell_counts)):
if cell_counts[k] < 196:
new_label[label_original == cell_idx[k]] = cell_idx[k]
# Do not exclude if the pixels are at the image boundaries.
_, W, H = label_original.shape
bd = np.zeros_like(label_original, dtype=label.dtype)
bd[:, 2 : W - 2, 2 : H - 2] = 1
new_label += label_original * bd
# Assign the transformed label
data["label"] = new_label
return data
class IntensityDiversification(MapTransform):
"""Randomly rescale the intensity of cell pixels."""
def __init__(
self,
keys=["img"],
change_cell_ratio=0.4,
scale_factors=[0, 0.7],
allow_missing_keys=False,
):
super(IntensityDiversification, self).__init__(keys, allow_missing_keys)
self.change_cell_ratio = change_cell_ratio
self.randscale_intensity = Compose(
[RandScaleIntensity(prob=1.0, factors=scale_factors)]
)
def __call__(self, data):
# Select cells to be transformed
cell_count = int(data["label"].max())
change_cell_count = int(cell_count * self.change_cell_ratio)
change_cell = np.random.choice(cell_count, change_cell_count, replace=False)
mask = copy.deepcopy(data["label"])
for i in range(cell_count):
cell_id = i + 1
if cell_id not in change_cell:
mask[mask == cell_id] = 0
mask[mask > 0] = 1
# Conduct intensity transformation for the selected cells
img_original = copy.deepcopy((1 - mask) * data["img"])
img_transformed = copy.deepcopy(mask * data["img"])
img_transformed = self.randscale_intensity(img_transformed)
# Assign the transformed image
data["img"] = img_original + img_transformed
return data
|