File size: 5,539 Bytes
fd4bbc8 |
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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
import cv2
import glob
import torch
import numpy as np
import albumentations as albu
from pathlib import Path
from albumentations.pytorch.transforms import ToTensorV2
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
class DatasetGenerate(Dataset):
def __init__(self, img_folder, gt_folder, edge_folder, phase: str = 'train', transform=None, seed=None):
self.images = sorted(glob.glob(img_folder + '/*'))
self.gts = sorted(glob.glob(gt_folder + '/*'))
self.edges = sorted(glob.glob(edge_folder + '/*'))
self.transform = transform
train_images, val_images, train_gts, val_gts, train_edges, val_edges = train_test_split(self.images, self.gts,
self.edges,
test_size=0.05,
random_state=seed)
if phase == 'train':
self.images = train_images
self.gts = train_gts
self.edges = train_edges
elif phase == 'val':
self.images = val_images
self.gts = val_gts
self.edges = val_edges
else: # Testset
pass
def __getitem__(self, idx):
image = cv2.imread(self.images[idx])
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask = cv2.imread(self.gts[idx])
mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
edge = cv2.imread(self.edges[idx])
edge = cv2.cvtColor(edge, cv2.COLOR_BGR2GRAY)
if self.transform is not None:
augmented = self.transform(image=image, masks=[mask, edge])
image = augmented['image']
mask = np.expand_dims(augmented['masks'][0], axis=0) # (1, H, W)
mask = mask / 255.0
edge = np.expand_dims(augmented['masks'][1], axis=0) # (1, H, W)
edge = edge / 255.0
return image, mask, edge
def __len__(self):
return len(self.images)
class Test_DatasetGenerate(Dataset):
def __init__(self, img_folder, gt_folder=None, transform=None):
self.images = sorted(glob.glob(img_folder + '/*'))
self.gts = sorted(glob.glob(gt_folder + '/*')) if gt_folder is not None else None
self.transform = transform
def __getitem__(self, idx):
image_name = Path(self.images[idx]).stem
image = cv2.imread(self.images[idx])
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
original_size = image.shape[:2]
if self.transform is not None:
augmented = self.transform(image=image)
image = augmented['image']
if self.gts is not None:
return image, self.gts[idx], original_size, image_name
else:
return image, original_size, image_name
def __len__(self):
return len(self.images)
def get_loader(img_folder, gt_folder, edge_folder, phase: str, batch_size, shuffle,
num_workers, transform, seed=None):
if phase == 'test':
dataset = Test_DatasetGenerate(img_folder, gt_folder, transform)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
else:
dataset = DatasetGenerate(img_folder, gt_folder, edge_folder, phase, transform, seed)
data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,
drop_last=True)
print(f'{phase} length : {len(dataset)}')
return data_loader
def get_train_augmentation(img_size, ver):
if ver == 1:
transforms = albu.Compose([
albu.Resize(img_size, img_size, always_apply=True),
albu.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
ToTensorV2(),
])
if ver == 2:
transforms = albu.Compose([
albu.OneOf([
albu.HorizontalFlip(),
albu.VerticalFlip(),
albu.RandomRotate90()
], p=0.5),
albu.OneOf([
albu.RandomContrast(),
albu.RandomGamma(),
albu.RandomBrightness(),
], p=0.5),
albu.OneOf([
albu.MotionBlur(blur_limit=5),
albu.MedianBlur(blur_limit=5),
albu.GaussianBlur(blur_limit=5),
albu.GaussNoise(var_limit=(5.0, 20.0)),
], p=0.5),
albu.Resize(img_size, img_size, always_apply=True),
albu.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
ToTensorV2(),
])
return transforms
def get_test_augmentation(img_size):
transforms = albu.Compose([
albu.Resize(img_size, img_size, always_apply=True),
albu.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225]),
ToTensorV2(),
])
return transforms
def gt_to_tensor(gt):
gt = cv2.imread(gt)
gt = cv2.cvtColor(gt, cv2.COLOR_BGR2GRAY) / 255.0
gt = np.where(gt > 0.5, 1.0, 0.0)
gt = torch.tensor(gt, device='cuda', dtype=torch.float32)
gt = gt.unsqueeze(0).unsqueeze(1)
return gt
|