|
|
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)
|
|
|
|
|
|
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_')))
|
|
|
]
|
|
|
|
|
|
|
|
|
random.shuffle(images)
|
|
|
|
|
|
|
|
|
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)) |
