Delete data_utils.py

data_utils.py

@@ -1,99 +0,0 @@
-import os
-
-import random
-import torch
-import imageio.v3 as imageio
-import numpy as np
-import skimage.morphology as morph
-import torchvision.transforms.v2.functional as T_F
-
-from skimage.filters import sato
-from pathlib import Path
-from scipy.ndimage import zoom
-from torchvision.datasets.folder import has_file_allowed_extension
-
-
-def make_dataset_t(image_dir, extensions=(".tif", ".tiff")):
-    image_dir = Path(image_dir)
-    images = [
-        (path, image_dir / f'Ridge_{path.name}', image_dir / f'Basins_{path.name}')
-        for path in sorted(image_dir.iterdir())
-        if (has_file_allowed_extension(path.name, extensions)
-            and (not path.name.startswith('Ridge_')) and (not path.name.startswith('Basins_')))
-    ]
-    return images
-
-def make_dataset_t_v(image_dir, extensions=(".tif", ".tiff")):
-    image_dir = Path(image_dir)
-    # Use list comprehension for faster filtering
-    images = [
-        (path, image_dir / f'Ridge_{path.name}', image_dir / f'Basins_{path.name}')
-        for path in sorted(image_dir.iterdir())
-        if (has_file_allowed_extension(path.name, extensions)
-            and (not path.name.startswith('Ridge_')) and (not path.name.startswith('Basins_')))
-    ]
-
-    # Shuffle in place
-    random.shuffle(images)
-
-    # Calculate split index once
-    split_idx = int(0.95 * len(images))
-    return images[:split_idx], images[split_idx:]
-
-def augmentations(image, label1, label2):
-    if random.random() < 0.5:
-        image, label1, label2 = T_F.vflip(image), T_F.vflip(label1), T_F.vflip(label2)
-    if random.random() < 0.5:
-        image, label1, label2 = T_F.hflip(image), T_F.hflip(label1), T_F.vflip(label2)
-    angles = [90, 180, 270]
-    angle = random.choice(angles)
-    if random.random() < 0.75:
-        image, label1, label2 = T_F.rotate(image, angle), T_F.rotate(label1, angle), T_F.rotate(label2, angle)
-    return image, label1, label2
-
-mean, std = (149.95293407563648, 330.8314960521203)
-target_water_level_range = [-100, 300]
-
-class TrainDataset(torch.utils.data.Dataset):
-    def __init__(self, train_split):
-        self.train_split = train_split
-
-    def __len__(self):
-        return len(self.train_split)
-
-    def __getitem__(self, index):
-        pair = self.train_split[index]
-        img = torch.from_numpy(imageio.imread(str(pair[0])))[None, :]
-        img = (img - mean) / std
-        ridge = torch.from_numpy(imageio.imread(str(pair[1])))[None, :].to(torch.float16)
-        basins = torch.from_numpy(imageio.imread(str(pair[2])))[None, :]
-        water_level = random.randint(*target_water_level_range)
-        basins = (basins >= water_level).to(torch.float16)
-        img, ridge, basins = augmentations(img, ridge, basins)
-        return img, ridge, basins, torch.tensor(water_level, dtype=torch.float16)
-
-class ValDataset(torch.utils.data.Dataset):
-    def __init__(self, val_split):
-        self.val_split = val_split
-
-    def __len__(self):
-        return len(self.val_split)
-
-    def __getitem__(self, index):
-        pair = self.val_split[index]
-        img = torch.from_numpy(imageio.imread(str(pair[0])))[None, :]
-        img = (img - mean) / std
-        ridge = torch.from_numpy(imageio.imread(str(pair[1])))[None, :].to(torch.float16)
-        basins = torch.from_numpy(imageio.imread(str(pair[2])))[None, :]
-        target_level = random.randint(*target_water_level_range)
-        basins = (basins >= target_level).to(torch.float16)
-        return img, ridge, basins, torch.tensor(target_level, dtype=torch.float16)
-
-if __name__ == '__main__':
-    train_split, val_split = make_dataset_t_v('dataset')
-
-    train_dataset = TrainDataset(train_split)
-    val_dataset = ValDataset(val_split)
-
-    print(train_dataset.__getitem__(0))
-    print(val_dataset.__getitem__(0))