| import os |
| import os.path |
| from typing import Any, Callable, cast, Dict, List, Optional, Tuple |
| from typing import Union |
|
|
| from PIL import Image |
| import pandas as pd |
| from torchvision.datasets import VisionDataset |
| import torch |
|
|
|
|
| def pil_loader(path: str) -> Image.Image: |
| |
| with open(path, "rb") as f: |
| img = Image.open(f) |
| return img.convert("RGB") |
|
|
| class BinaryWaterbirds(VisionDataset): |
| def __init__( |
| self, |
| root: str, |
| split: str, |
| loader: Callable[[str], Any] = pil_loader, |
| transform: Optional[Callable] = None, |
| target_transform: Optional[Callable] = None, |
| ) -> None: |
| super().__init__(root, transform=transform, target_transform=target_transform) |
| |
| self.loader = loader |
| csv = pd.read_csv(os.path.join(root, 'metadata.csv')) |
| split = {'test': 2, 'valid': 1, 'train': 0}[split] |
| csv = csv[csv['split'] == split] |
| self.samples = [(os.path.join(root, csv.iloc[i]['img_filename']), csv.iloc[i]['y']) for i in range(len(csv))] |
| |
| def __getitem__(self, index: int) -> Tuple[Any, Any]: |
| """ |
| Args: |
| index (int): Index |
| Returns: |
| tuple: (sample, target) where target is class_index of the target class. |
| """ |
| path, target = self.samples[index] |
| sample = self.loader(path) |
| if self.transform is not None: |
| sample = self.transform(sample) |
| if self.target_transform is not None: |
| target = self.target_transform(target) |
|
|
| return sample, target |
|
|
| def __len__(self) -> int: |
| return len(self.samples) |
|
|
