garmentiq.classification.load_data

 1from torchvision import transforms
 2import os
 3from PIL import Image
 4import torch
 5from garmentiq.classification.utils import (
 6    CachedDataset,
 7    seed_worker,
 8    train_epoch,
 9    validate_epoch,
10    save_best_model,
11    validate_train_param,
12    validate_test_param,
13)
14from tqdm.notebook import tqdm
15
16
17def load_data(
18    df,
19    img_dir,
20    label_column,
21    resize_dim=(120, 184),
22    normalize_mean=[0.8047, 0.7808, 0.7769],
23    normalize_std=[0.2957, 0.3077, 0.3081],
24):
25    """
26    Loads and preprocesses image data into memory from a DataFrame of filenames and labels.
27
28    This function reads images from the specified directory, applies resizing, normalization,
29    and tensor conversion, and encodes labels from a specified column. It returns tensors for
30    images and labels, along with the transform pipeline used.
31
32    Args:
33        df (pandas.DataFrame): A pandas DataFrame containing at least a 'filename' column and a label column.
34        img_dir (str): Path to the directory containing image files.
35        label_column (str): Name of the column in `df` containing class labels.
36        resize_dim (tuple[int, int]): Tuple indicating the dimensions (height, width) to resize each image to.
37                                     Defaults to (120, 184).
38        normalize_mean (list[float]): Mean values for normalization (per channel).
39                                      Defaults to `[0.8047, 0.7808, 0.7769]`.
40        normalize_std (list[float]): Standard deviation values for normalization (per channel).
41                                     Defaults to `[0.2957, 0.3077, 0.3081]`.
42
43    Returns:
44        tuple[torch.Tensor, torch.Tensor, torchvision.transforms.Compose]: A tuple containing:
45            - cached_images (torch.Tensor): Tensor containing all preprocessed images.
46            - cached_labels (torch.Tensor): Tensor containing all encoded labels.
47            - transform (torchvision.transforms.Compose): The transformation pipeline used.
48    """
49    transform = transforms.Compose(
50        [
51            transforms.Resize(resize_dim),
52            transforms.ToTensor(),
53            transforms.Normalize(mean=normalize_mean, std=normalize_std),
54        ]
55    )
56
57    classes = sorted(df[label_column].unique())
58    class_to_idx = {c: i for i, c in enumerate(classes)}
59
60    cached_images = []
61    cached_labels = []
62
63    for _, row in tqdm(df.iterrows(), total=len(df), desc="Loading data into memory"):
64        img_path = os.path.join(img_dir, row["filename"])
65        image = Image.open(img_path).convert("RGB")
66        image = transform(image)
67
68        label = class_to_idx[row[label_column]]
69
70        cached_images.append(image)
71        cached_labels.append(label)
72
73    cached_images = torch.stack(cached_images)
74    cached_labels = torch.tensor(cached_labels)
75
76    return cached_images, cached_labels, transform