| import argparse |
| import os |
| import random |
| import numpy as np |
| import torch |
| import cv2 |
| import tifffile |
| import matplotlib.pyplot as plt |
| from typing import List, Tuple |
| from torch.utils.data import Dataset |
| from torchvision.transforms import v2 |
| from torchvision import tv_tensors |
|
|
| import sys |
|
|
| |
| sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| from affmae.data.transforms import ElasticTransform |
|
|
|
|
| IMAGE_MEAN = [0.0] |
| IMAGE_STD = [1.0] |
| THRESHOLD = 0.5 |
|
|
|
|
| class EMDatasetMultiClass(Dataset): |
| def __init__( |
| self, |
| base_path: str, |
| test_dataset: bool = False, |
| img_size: int = 512, |
| apply_transforms: bool = False, |
| indices: List[int] = [0, 1, 2], |
| input_ext: str = ".tif", |
| return_path: bool = True, |
| ): |
| self.test_dataset = test_dataset |
| self.img_size = img_size |
| self.apply_transforms = apply_transforms |
| self.indices = indices |
| self.return_path = return_path |
|
|
| |
| train_place_in = os.path.join(base_path, "train") |
| test_place_in = os.path.join(base_path, "test") |
|
|
| if test_dataset: |
| input_folder = os.path.join(test_place_in, "images") |
| output_folder = os.path.join(test_place_in, "masks") |
| else: |
| input_folder = os.path.join(train_place_in, "images") |
| output_folder = os.path.join(train_place_in, "masks") |
| |
| if not os.path.exists(input_folder): |
| raise FileNotFoundError(f"Input folder not found: {input_folder}") |
|
|
| self.image_paths, self.mask_paths = self._load_image_names(input_folder, output_folder, input_ext) |
|
|
| self.img_mean = torch.tensor(IMAGE_MEAN, dtype=torch.float) |
| self.img_std = torch.tensor(IMAGE_STD, dtype=torch.float) |
|
|
| self._setup_transforms() |
|
|
| @staticmethod |
| def _load_image_names(input_folder: str, output_folder: str, input_ext: str = ".tif") -> Tuple[np.ndarray, np.ndarray]: |
| all_img = np.array(os.listdir(input_folder)) |
| all_tiff = np.array(os.listdir(output_folder)) |
| split_func = np.vectorize(lambda x: x.split(".")[0]) |
| |
| |
| img_paths = np.intersect1d(split_func(all_img), split_func(all_tiff)) |
|
|
| append_input = np.vectorize(lambda x: os.path.join(input_folder, x + input_ext)) |
| input_imgs = append_input(img_paths) |
|
|
| append_output = np.vectorize(lambda x: os.path.join(output_folder, x + ".tiff")) |
| output_imgs = append_output(img_paths) |
|
|
| return input_imgs, output_imgs |
|
|
| def _setup_transforms(self): |
| base_transforms = [ |
| v2.ToImage(), |
| v2.ConvertImageDtype(torch.float32), |
| v2.Resize(size=(self.img_size, self.img_size), antialias=True, interpolation=v2.InterpolationMode.BILINEAR), |
| ] |
|
|
| if self.test_dataset or not self.apply_transforms: |
| self.transforms = v2.Compose(base_transforms) |
| else: |
| augmentation_transforms = [ |
| v2.RandomHorizontalFlip(p=0.5), |
| v2.RandomVerticalFlip(p=0.5), |
| v2.RandomAdjustSharpness(sharpness_factor=0.8, p=0.25), |
| v2.RandomAdjustSharpness(sharpness_factor=1.25, p=0.25), |
| v2.RandomAffine( |
| degrees=30, |
| translate=(0.0, 0.1), |
| scale=(0.9, 1.3), |
| shear=20.0 |
| ), |
| ElasticTransform(alpha=(100, 100), sigma=(10, 10), p=0.6) |
| ] |
| self.transforms = v2.Compose(base_transforms + augmentation_transforms) |
|
|
| def __len__(self) -> int: |
| return len(self.image_paths) |
|
|
| def _load_segm_image(self, id: int) -> np.ndarray: |
| img = tifffile.imread(self.image_paths[id]) |
| return img if len(img.shape) == 3 else img.reshape(img.shape + (1,)) |
|
|
| def _load_segm_target(self, id: int) -> np.ndarray: |
| img = tifffile.imread(self.mask_paths[id]) |
| if len(img.shape) == 4: img = img[:, :, :, 0] |
| return img if len(img.shape) == 2 else img[[self.indices] if isinstance(self.indices, int) else self.indices, :, :] |
|
|
| def __getitem__(self, index: int): |
| image_np = self._load_segm_image(index) |
| target_np_multi_channel = self._load_segm_target(index) |
|
|
| |
| clahe = cv2.createCLAHE(clipLimit=4.25, tileGridSize=(8, 8)) |
| try: |
| if image_np.dtype != np.uint8: |
| norm_img = cv2.normalize(image_np, None, 0, 255, cv2.NORM_MINMAX, dtype=cv2.CV_8U) |
| image_np = clahe.apply(norm_img) |
| else: |
| image_np = clahe.apply(image_np) |
| except: |
| if len(image_np.shape) == 3 and image_np.shape[2] == 3: |
| image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2GRAY) |
| image_np = clahe.apply(image_np) |
|
|
| if len(image_np.shape) == 2: |
| image_np = image_np.reshape(image_np.shape + (1,)) |
|
|
| if len(target_np_multi_channel.shape) == 3: |
| C, H, W = target_np_multi_channel.shape |
| target_np_multi_class = np.zeros((H, W), dtype=np.int64) |
| for c in range(C): |
| class_id = c + 1 |
| is_active = target_np_multi_channel[c, :, :] > THRESHOLD |
| target_np_multi_class[is_active] = class_id |
| else: |
| target_np_multi_class = target_np_multi_channel |
|
|
| image = image_np |
| target = tv_tensors.Mask(target_np_multi_class) |
|
|
| image, target = self.transforms(image, target) |
| image = image.float().sub_(self.img_mean).div_(self.img_std) |
|
|
| if self.return_path: |
| return image, target, (self.image_paths[index], self.mask_paths[index]) |
| return image, target |
|
|
|
|
| def parse_args(): |
| parser = argparse.ArgumentParser( |
| description="Render a grid of augmented training samples to eyeball the " |
| "augmentation pipeline.") |
| parser.add_argument("--base-path", required=True, |
| help="Dataset root (the directory holding train/test).") |
| parser.add_argument("--output", default="aug_viz.png", |
| help="Output image path.") |
| parser.add_argument("--img-size", type=int, default=512) |
| parser.add_argument("--input-ext", default=".tiff", |
| help="Image file extension.") |
| parser.add_argument("--config", default=None, |
| help="Optional YAML config to take base_path, img_size and " |
| "input_ext from instead of the flags above.") |
| return parser.parse_args() |
|
|
|
|
| def visualize_grid(base_path, output_path, img_size, input_ext): |
| print(f"Loading dataset from: {base_path}") |
| try: |
| dataset = EMDatasetMultiClass( |
| base_path=base_path, |
| test_dataset=False, |
| img_size=img_size, |
| apply_transforms=True, |
| input_ext=input_ext, |
| return_path=True |
| ) |
| except Exception as e: |
| print(f"Error initializing dataset: {e}") |
| return |
|
|
| print(f"Found {len(dataset)} images.") |
| |
| if len(dataset) < 25: |
| print("Dataset is too small for a 5x5 grid, using all images.") |
| indices = list(range(len(dataset))) |
| else: |
| indices = random.sample(range(len(dataset)), 25) |
|
|
| fig, axes = plt.subplots(5, 5, figsize=(15, 15)) |
| fig.suptitle("5x5 Random Augmentation Grid", fontsize=16) |
|
|
| print("Generating grid...") |
| |
| for idx, ax in zip(indices, axes.flat): |
| img_tensor, mask_tensor, (img_path, mask_path) = dataset[idx] |
| |
| rel_path = os.path.relpath(img_path, base_path) |
| |
| |
| img = img_tensor.cpu().numpy() |
| img = np.transpose(img, (1, 2, 0)) |
| mean = np.array(IMAGE_MEAN) |
| std = np.array(IMAGE_STD) |
| img = (img * std) + mean |
| img = np.clip(img, 0.0, 1.0) |
| |
| if img.shape[2] == 1: |
| img = img.squeeze(2) |
| ax.imshow(img, cmap='gray') |
| else: |
| ax.imshow(img) |
| |
| |
| ax.set_title(rel_path, fontsize=6) |
| ax.axis('off') |
|
|
| plt.tight_layout() |
| out_dir = os.path.dirname(os.path.abspath(output_path)) |
| os.makedirs(out_dir, exist_ok=True) |
|
|
| plt.savefig(output_path, dpi=125) |
| print(f"Saved visualization to: {output_path}") |
|
|
|
|
| def main(): |
| args = parse_args() |
| base_path, img_size, input_ext = args.base_path, args.img_size, args.input_ext |
| if args.config: |
| from affmae.config import load_config |
| cfg = load_config(args.config) |
| base_path = cfg.base_path |
| img_size = cfg.img_size |
| input_ext = cfg.input_ext |
| visualize_grid(base_path, args.output, img_size, input_ext) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
