File size: 2,769 Bytes
a2aee5d | 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 | """
dataset.py — PyTorch datasets wrapping the BDAPPV HuggingFace dataset.
Two dataset classes:
- SegmentationDataset : positives only (has_mask=True), returns (image, mask)
- ClassificationDataset : all examples, returns (image, label)
"""
import torch
from torch.utils.data import Dataset
from torchvision import transforms as T
IMG_MEAN = [0.485, 0.456, 0.406]
IMG_STD = [0.229, 0.224, 0.225]
class SegmentationDataset(Dataset):
"""
Wraps a HF split pre-filtered to positives (has_mask=True).
Returns (image_tensor [3, H, W], mask_tensor [1, H, W]) as float32.
"""
def __init__(self, hf_split, img_size: int = 400, augment: bool = False):
self.data = hf_split.filter(lambda x: x["has_mask"])
self.augment = augment
self.img_size = img_size
self.img_tf = T.Compose([
T.Resize((img_size, img_size)),
T.ToTensor(),
T.Normalize(mean=IMG_MEAN, std=IMG_STD),
])
self.mask_tf = T.Compose([
T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
T.ToTensor(),
])
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
ex = self.data[idx]
image = ex["image"].convert("RGB")
mask = ex["mask"].convert("L")
if self.augment:
if torch.rand(1) > 0.5:
image = T.functional.hflip(image)
mask = T.functional.hflip(mask)
if torch.rand(1) > 0.5:
image = T.functional.vflip(image)
mask = T.functional.vflip(mask)
image = self.img_tf(image)
mask = self.mask_tf(mask)
mask = (mask > 0.5).float() # binarize → {0, 1}
return image, mask
class ClassificationDataset(Dataset):
"""
All examples (positives + negatives).
Returns (image_tensor [3, H, W], label [float scalar]).
label = 1.0 if has_mask else 0.0
"""
def __init__(self, hf_split, img_size: int = 299, augment: bool = False):
self.data = hf_split
self.augment = augment
self.img_tf = T.Compose([
T.Resize((img_size, img_size)),
T.ToTensor(),
T.Normalize(mean=IMG_MEAN, std=IMG_STD),
])
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
ex = self.data[idx]
image = ex["image"].convert("RGB")
label = float(ex["has_mask"])
if self.augment:
if torch.rand(1) > 0.5:
image = T.functional.hflip(image)
if torch.rand(1) > 0.5:
image = T.functional.vflip(image)
image = self.img_tf(image)
return image, torch.tensor(label)
|