fairfacedata.py

import os
import pickle
from typing import Any

import numpy as np
import pandas as pd
import torch
from PIL import Image
from pytorch_lightning import LightningDataModule
from torch.utils.data import (DataLoader, Dataset, Subset, TensorDataset,
                              random_split)
from torchvision import transforms

from utils import add_label_bias, count_label_flips

ItemType = tuple[Any, Any, Any, int]


class FairFaceDataset(Dataset):
    def __init__(self, csv_path, img_root, transform=None):
        self.df = pd.read_csv(csv_path)
        self.df = self.df[self.df["race"].isin(["White", "Black"])]
        self.img_root = img_root
        self.transform = transform
        self.df["sens"] = self.df["race"].apply(lambda r: 1 if r == "Black" else 0)
        self.df["label"] = self.df["gender"].map({"Male": 1, "Female": 0})
        self.df = self.df.reset_index(drop=True)

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        idx = int(idx) 
        row = self.df.iloc[idx]
        img_path = os.path.join(self.img_root, row["file"])
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, row["label"], row["sens"], idx


class FairFaceDataModule(LightningDataModule):
    def __init__(self, 
                 data_location, 
                 batch_size=128,
                 bias_amount=0.4,
                 meta_ratio=0.001,
                 num_workers=0,
                 train_val_split=0.8,
                 debug=False): 
        super().__init__()
        self.train_csv = os.path.join(data_location, "fairface_label_train.csv")
        self.val_csv = os.path.join(data_location, "fairface_label_val.csv")
        self.img_dir = os.path.join(data_location, "fairface-img-margin025-trainval")
        self.batch_size = batch_size
        self.bias_amount = bias_amount
        self.meta_ratio = meta_ratio
        self.num_workers = num_workers
        self.train_val_split = train_val_split
        self.debug = debug
    
        self.theta_dict = {
            'theta_0_p': bias_amount,
            'theta_0_m': 0,
            'theta_1_p': 0,
            'theta_1_m': bias_amount
        }

        self.train_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])
        ]) 

        self.val_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])
        ])

    def setup(self, stage=None):
        full_train_dataset = FairFaceDataset(self.train_csv, self.img_dir, transform=None)
        all_data = []
        for idx in range(len(full_train_dataset)):
            x, y, s, _ = full_train_dataset[idx]
            all_data.append((x, y, s, idx))

        images = [item[0] for item in all_data]
        y_labels = torch.tensor([item[1] for item in all_data], dtype=torch.long)  # gender
        s_labels = torch.tensor([item[2] for item in all_data], dtype=torch.long)  # race
        indices = torch.tensor([item[3] for item in all_data], dtype=torch.long)

        combined_labels = y_labels * 2 + s_labels  # (0,0), (0,1), (1,0), (1,1)
        unique_combinations = torch.unique(combined_labels)

        train_indices = []
        test_indices = []
        np.random.seed(42)

        for combo in unique_combinations:
            combo_mask = (combined_labels == combo)
            combo_indices = torch.where(combo_mask)[0].numpy()
            if len(combo_indices) > 0:
                np.random.shuffle(combo_indices)
                
                if self.debug:
                    combo_indices = combo_indices[:min(100, len(combo_indices))]

                n_total = len(combo_indices)
                n_test = max(1, int(n_total * (1 - self.train_val_split)))
                n_train = n_total - n_test

                if n_train < 1:
                    n_train = max(1, n_total - 1)
                    n_test = n_total - n_train

                test_indices.extend(combo_indices[:n_test])
                train_indices.extend(combo_indices[n_test:])

        train_indices = np.array(train_indices)
        test_indices = np.array(test_indices)

        np.random.shuffle(train_indices)
        np.random.shuffle(test_indices)

        n_meta = max(1, int(self.meta_ratio * len(train_indices)))
        meta_indices = train_indices[:n_meta]
        train_indices = train_indices[n_meta:]

        y_train = y_labels[train_indices]
        s_train = s_labels[train_indices]
        y_meta = y_labels[meta_indices]
        s_meta = s_labels[meta_indices]
        y_test = y_labels[test_indices]
        s_test = s_labels[test_indices]

        if self.bias_amount != 0:
            print("BIAS ADDED")
            y_train_biased = add_label_bias(y_train, s_train, self.theta_dict)
        else:
            print("ASSUME NO BIAS")
            y_train_biased = y_train

        count_label_flips(y_train, y_train_biased)
        train_images = torch.stack([self.train_transform(images[i]) for i in train_indices])

        self.train_dataset = TensorDataset(
            train_images,
            y_train_biased, 
            s_train,
            torch.tensor(range(len(train_indices)), dtype=torch.long)  
        )

        meta_images = torch.stack([self.val_transform(images[i]) for i in meta_indices])
        self.meta_dataset = TensorDataset(
            meta_images,
            y_meta,
            s_meta,
            torch.tensor(range(len(meta_indices)), dtype=torch.long) 
        )


        test_images = torch.stack([self.val_transform(images[i]) for i in test_indices])
        self.val_dataset = TensorDataset(
            test_images,
            y_test,
            s_test,
            torch.tensor(range(len(test_indices)), dtype=torch.long)  
        )


        unique_groups, group_counts = torch.unique(s_train, return_counts=True)
        total_train = len(s_train)
        self.global_group_distribution = {int(g.item()): c.item() / total_train for g, c in zip(unique_groups, group_counts)}

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size,
                          shuffle=True, num_workers=self.num_workers)

    def val_dataloader(self):
        return DataLoader(self.val_dataset, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers)

    def test_dataloader(self):
        return self.val_dataloader()


class DualLoader:
    def __init__(self, train_loader, meta_dataset):
        self.train_loader = train_loader
        self.meta_dataset = meta_dataset

        x_m_list, y_m_list, s_m_list, idx_m_list = [], [], [], []
        for i in range(len(self.meta_dataset)):
            x, y, s, idx = self.meta_dataset[i]
            x_m_list.append(x)
            y_m_list.append(y)
            s_m_list.append(s)
            idx_m_list.append(idx)

        self.x_m = torch.stack(x_m_list)
        self.y_m = torch.tensor(y_m_list)
        self.s_m = torch.tensor(s_m_list)
        self.idx_m = torch.tensor(idx_m_list)

    def __iter__(self):
        self.train_iter = iter(self.train_loader)
        return self

    def __next__(self):
        try:
            train_batch = next(self.train_iter)
        except StopIteration:
            raise StopIteration

        x_t, y_t, s_t, idx_t = train_batch
        return {
            'train': (x_t, y_t, s_t, idx_t), 
            'meta': (self.x_m, self.y_m, self.s_m, self.idx_m)
        }
    
    def __len__(self):
        return len(self.train_loader)