Spaces:
Paused
Paused
| from torch.utils.data import DataLoader | |
| import cv2 | |
| import numpy as np | |
| import albumentations as A | |
| from albumentations.pytorch import ToTensorV2 | |
| from datasets import load_dataset | |
| from torch.utils.data import Dataset | |
| from src.operations import remap_mask | |
| class ClothesDataset(Dataset): | |
| def __init__(self, dataset, transform = None, label_map = None): | |
| self.dataset = dataset | |
| self.transform = transform | |
| self.label_map = label_map | |
| def __len__(self): | |
| return len(self.dataset) | |
| def __getitem__(self, idx): | |
| sample = self.dataset[idx] | |
| # Image | |
| image = sample["pixel_values"].convert("RGB") | |
| image = np.array(image, dtype=np.uint8) | |
| # Mask | |
| mask = np.array(sample["label"], dtype = np.uint8) | |
| # Remap Mask | |
| mask = remap_mask(mask, self.label_map) | |
| # Apply transform | |
| if self.transform: | |
| transformed = self.transform(image=image, mask=mask) | |
| image = transformed["image"] | |
| mask = transformed["mask"] | |
| image = image.contiguous() | |
| mask = mask.contiguous() | |
| return image, mask.long() | |
| def get_transforms(config, train=True): | |
| image_size = config["data"]["image_size"] | |
| mean = config["data"]["normalization"]["mean"] | |
| std = config["data"]["normalization"]["std"] | |
| transforms = [ | |
| A.LongestMaxSize( | |
| max_size = image_size, | |
| interpolation=cv2.INTER_LINEAR, | |
| mask_interpolation=cv2.INTER_NEAREST | |
| ), | |
| A.PadIfNeeded( | |
| min_height=image_size, | |
| min_width=image_size, | |
| border_mode=cv2.BORDER_CONSTANT, | |
| fill = 0, | |
| fill_mask = 0 | |
| ) | |
| ] | |
| if train: | |
| transforms.append( | |
| A.HorizontalFlip( | |
| p = config["data"]["augmentation"]["horizontal_flip_probability"] | |
| ) | |
| ) | |
| transforms.extend([ | |
| A.Normalize( | |
| mean = mean, | |
| std=std, | |
| ), | |
| ToTensorV2() | |
| ]) | |
| return A.Compose(transforms) | |
| def create_dataloaders(config): | |
| # load dataset | |
| dataset = load_dataset("ckotait/ATRDataset") | |
| # Transforms | |
| train_transform = get_transforms(config, train=True) | |
| val_transform = get_transforms(config, train=False) | |
| # Create dataset class | |
| train_dataset = ClothesDataset(dataset['train'], train_transform, config["segmentation"]["label_map"]) | |
| val_dataset = ClothesDataset(dataset['validation'], val_transform, config["segmentation"]["label_map"]) | |
| test_dataset = ClothesDataset(dataset['test'], val_transform, config["segmentation"]["label_map"]) | |
| # Data loaders | |
| train_loader = DataLoader( | |
| train_dataset, | |
| batch_size=config["data"]["batch_size"], | |
| shuffle=True, | |
| num_workers=config["data"]["num_workers"], | |
| pin_memory=True, | |
| drop_last=True | |
| ) | |
| val_loader = DataLoader( | |
| val_dataset, | |
| batch_size=config["data"]["batch_size"], | |
| shuffle=False, | |
| num_workers=config["data"]["num_workers"], | |
| pin_memory=True | |
| ) | |
| test_loader = DataLoader( | |
| test_dataset, | |
| batch_size=config["data"]["batch_size"], | |
| shuffle=False, | |
| num_workers=config["data"]["num_workers"], | |
| pin_memory=True | |
| ) | |
| return train_loader, val_loader, test_loader | |