garmentiq.classification.load_data

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