Skip to content

BasicDataset

anomaly_match.datasets.BasicDataset

Basic PyTorch dataset wrapper for image loading and augmentation.

BasicDataset(data, filenames, targets=None, num_classes=None, transform=None, use_strong_transform=False, strong_transform=None, num_channels=3)

Bases: Dataset

BasicDataset returns a pair of image and labels (targets).

If targets are not given, BasicDataset returns None as the label. This class supports strong augmentation for Fixmatch, and returns both weakly and strongly augmented images.

Parameters:

Name Type Description Default
data ndarray | list | Tensor

x_data as numpy array or torch tensor

 required
filenames list[str]

filenames of x_data

 required
targets Tensor | list | None

y_data (if not exist, None)

 None
num_classes int | None

number of label classes

 None
transform Callable | None

basic transformation of data

 None
use_strong_transform bool

If True, this dataset returns both weakly and strongly augmented images.

 False
strong_transform Callable | None

list of transformation functions for strong augmentation

 None
num_channels int

number of image channels

 3
Source code in anomaly_match/datasets/BasicDataset.py 
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
def __init__(
    self,
    data: np.ndarray | list | torch.Tensor,
    filenames: list[str],
    targets: torch.Tensor | list | None = None,
    num_classes: int | None = None,
    transform: Callable | None = None,
    use_strong_transform: bool = False,
    strong_transform: Callable | None = None,
    num_channels: int = 3,
) -> None:
    super(BasicDataset, self).__init__()

    # Ensure data and filenames are of the same length
    assert len(data) == len(filenames), "data and filenames must have the same length"
    if targets is not None:
        assert len(data) == len(targets), "data and targets must have the same length"

    # Convert data to shared memory to minimize duplication across workers
    if isinstance(data, np.ndarray):
        self.data = torch.from_numpy(data)  # Convert to torch tensor
    elif isinstance(data, list):
        if len(data) > 0:
            self.data = torch.stack([torch.tensor(d) for d in data])
        else:
            self.data = torch.empty(0)  # Handle the empty dataset case
    else:
        self.data = data  # Assume it's a torch tensor already

    # If data is non-empty, place tensor in shared memory to reduce duplication
    if len(self.data) > 0:
        self.data.share_memory_()

    self.filenames = filenames
    if targets is not None:
        if isinstance(targets, list):
            targets = torch.tensor(targets)
        self.targets = targets.clone().detach()

    self.num_classes = num_classes
    self.num_channels = num_channels
    self.use_strong_transform = use_strong_transform

    self.transform = transform
    if use_strong_transform:
        if strong_transform is None:
            self.strong_transform = get_strong_transforms(num_channels=num_channels)
        else:
            self.strong_transform = strong_transform

__getitem__(idx)

Return augmented image, target, and filename.

If strong augmentation is not used, return weak_augment_image, target, filename. Otherwise return weak_augment_image, strong_augment_image, target.

Source code in anomaly_match/datasets/BasicDataset.py 
 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
def __getitem__(self, idx):
    """Return augmented image, target, and filename.

    If strong augmentation is not used,
    return weak_augment_image, target, filename.
    Otherwise return weak_augment_image, strong_augment_image, target.
    """
    # Set idx-th target
    target = self.targets[idx] if self.targets is not None else None

    # Set augmented images
    img = self.data[idx]
    if self.transform is None:
        return transforms.ToTensor()(img), target, self.filenames[idx]
    else:
        # Convert to numpy if tensor
        if isinstance(img, torch.Tensor):
            img_np = img.numpy()
        else:
            img_np = img

        # For RGB images (3 channels), use PIL-based transforms
        # For N-channel images, pass numpy array directly to transforms
        if self.num_channels == 3:
            img = Image.fromarray(img_np)
        else:
            img = img_np

        img_w = self.transform(img)

    if not self.use_strong_transform:
        return img_w, target, self.filenames[idx]
    else:
        img_s = self.strong_transform(img)
        return img_w, img_s, target