Spaces:
No application file
No application file
| from torch.utils.data import Dataset, DataLoader, random_split | |
| import torch | |
| from segment_anything.utils.transforms import ResizeLongestSide | |
| from .data_fetching import DataFetcher | |
| from .data_utils import DataProcessing | |
| class HistologyDataset(Dataset): | |
| def __init__(self, datasets, data_directory, cluster, image_encoder_size, mask_augmentation_tries, data_augmentations, threshold_connected_components): | |
| self.datasets = datasets | |
| self.data_fetcher = DataFetcher(data_directory, cluster) | |
| self.image_encoder_size = image_encoder_size | |
| self.mask_augmentation_tries = mask_augmentation_tries | |
| self.data_augmentations = data_augmentations | |
| self.threshold_connected_components = threshold_connected_components | |
| self.pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1) | |
| self.pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1) | |
| self.DATASET_DICT = { | |
| "BCSS": [self.data_fetcher.len_bcss, self.data_fetcher.get_bcss, False, 0.0], | |
| "CAMELYON": [self.data_fetcher.len_camelyon, self.data_fetcher.get_camelyon, True, 0.0], | |
| "CellSeg": [self.data_fetcher.len_cellseg, self.data_fetcher.get_cellseg, False, 0.0], | |
| "CoCaHis": [self.data_fetcher.len_cocahis, self.data_fetcher.get_cocahis, False, 0.0], | |
| "CoNIC": [self.data_fetcher.len_conic, self.data_fetcher.get_conic, False, 0.0], | |
| "CPM": [self.data_fetcher.len_cpm, self.data_fetcher.get_cpm, False, 0.0], | |
| "CRAG": [self.data_fetcher.len_crag, self.data_fetcher.get_crag, False, 1.0], | |
| "CryoNuSeg": [self.data_fetcher.len_cryonuseg, self.data_fetcher.get_cryonuseg, False, 0.0], | |
| "GlaS": [self.data_fetcher.len_glas, self.data_fetcher.get_glas, False, 1.0], | |
| "ICIA2018": [self.data_fetcher.len_icia2018, self.data_fetcher.get_icia2018, True, 0.0], | |
| "Janowczyk": [self.data_fetcher.len_janowczyk, self.data_fetcher.get_janowczyk, False, 0.0], | |
| "KPI": [self.data_fetcher.len_kpi, self.data_fetcher.get_kpi, False, 0.0], | |
| "Kumar": [self.data_fetcher.len_kumar, self.data_fetcher.get_kumar, False, 0.0], | |
| "MoNuSAC": [self.data_fetcher.len_monusac, self.data_fetcher.get_monusac, False, 0.0], | |
| "MoNuSeg": [self.data_fetcher.len_monuseg, self.data_fetcher.get_monuseg, False, 0.0], | |
| "NuClick": [self.data_fetcher.len_nuclick, self.data_fetcher.get_nuclick, False, 0.0], | |
| "PAIP2023": [self.data_fetcher.len_paip2023, self.data_fetcher.get_paip2023, False, 0.0], | |
| "PanNuke": [self.data_fetcher.len_pannuke, self.data_fetcher.get_pannuke, False, 0.0], | |
| "SegPath": [self.data_fetcher.len_segpath, self.data_fetcher.get_segpath, False, 0.0], | |
| "SegPC": [self.data_fetcher.len_segpc, self.data_fetcher.get_segpc, False, 0.0], | |
| "TIGER": [self.data_fetcher.len_tiger, self.data_fetcher.get_tiger, False, 0.0], | |
| "TNBC": [self.data_fetcher.len_tnbc, self.data_fetcher.get_tnbc, False, 0.0], | |
| } | |
| self.len = 0 | |
| for dataset in self.datasets: | |
| self.DATASET_DICT[dataset][0] = self.DATASET_DICT[dataset][0]() | |
| self.len += self.DATASET_DICT[dataset][0] | |
| self.resize = ResizeLongestSide(image_encoder_size) | |
| def __len__(self): | |
| return self.len | |
| def __getitem__(self, idx): | |
| for dataset in self.datasets: | |
| if idx < self.DATASET_DICT[dataset][0]: | |
| image_name, mask_name = self.DATASET_DICT[dataset][1](idx) | |
| return idx, DataProcessing.preprocess( | |
| image_name, mask_name, self.pixel_mean, self.pixel_std, self.DATASET_DICT[dataset][2], | |
| self.image_encoder_size, self.mask_augmentation_tries, self.data_augmentations, self.threshold_connected_components | |
| ), self.DATASET_DICT[dataset][3] | |
| else: | |
| idx -= self.DATASET_DICT[dataset][0] | |
| def get_image(self, idx): | |
| for dataset in self.datasets: | |
| if idx < self.DATASET_DICT[dataset][0]: | |
| image_name, mask_name = self.DATASET_DICT[dataset][1](idx) | |
| return image_name, mask_name | |
| else: | |
| idx -= self.DATASET_DICT[dataset][0] | |
| def prepare_data(data_config): | |
| """ | |
| """ | |
| datasets = data_config["datasets"] | |
| data_directory = data_config["data_directory"] | |
| cluster = data_config["cluster"] | |
| image_encoder_size = data_config["image_encoder_size"] | |
| batch_size = data_config["batch_size"] | |
| drop_last = data_config["drop_last"] | |
| num_workers = data_config["num_workers"] | |
| # Parameters only for training | |
| use_holdout_testset = data_config.get("use_holdout_testset", True) | |
| holdout_testsets = data_config.get("test_datasets", datasets) | |
| train_split = data_config.get("train_split", 0.8) | |
| val_split = data_config.get("val_split", 0.1) | |
| test_split = data_config.get("test_split", 0.1) | |
| data_split = data_config.get("data_split", False) | |
| shuffle = data_config.get("shuffle", False) | |
| seed = data_config.get("seed", 1) | |
| mask_augmentation_tries = data_config.get("mask_augmentation_tries", 5) | |
| data_augmentations = data_config.get("data_augmentations", ["NoOp"]) | |
| threshold_connected_components = data_config.get("threshold_connected_components", 2) | |
| dataset = HistologyDataset(datasets, data_directory, cluster, image_encoder_size, mask_augmentation_tries, data_augmentations, threshold_connected_components) | |
| if data_split: | |
| generator = torch.Generator().manual_seed(seed) | |
| if use_holdout_testset: | |
| train_set, val_set = random_split(dataset, [train_split + val_split, test_split],generator=generator) | |
| test_set = HistologyDataset(holdout_testsets, data_directory, cluster, image_encoder_size, mask_augmentation_tries, ["NoOp"], threshold_connected_components) | |
| else: | |
| train_set, val_set, test_set = random_split(dataset, [train_split, val_split, test_split],generator=generator) | |
| train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers) | |
| val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, drop_last=drop_last, num_workers=num_workers) | |
| test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, drop_last=drop_last, num_workers=num_workers) | |
| return train_loader, val_loader, test_loader | |
| else: | |
| dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=drop_last, num_workers=num_workers) | |
| return dataset, dataloader |