File size: 4,404 Bytes
bd659a9 | 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 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 | from pathlib import Path
from typing import Dict, List, Tuple
import cv2
import pandas as pd
import torch
from torch.utils.data import Dataset
def normalize_label(value: object) -> float:
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"misogyny", "misogynous", "1", "true"}:
return 1.0
if normalized in {"not-misogyny", "not misogyny", "non-misogyny", "0", "false"}:
return 0.0
raise ValueError(f"Unsupported label value: {value!r}")
return float(value)
class MemeDataset(Dataset):
def __init__(
self,
split: str,
data_root: Path,
language: str,
tokenizer,
image_id_col: str,
text_col: str,
label_col: str,
max_length: int,
oversample_positive_train: int = 0,
) -> None:
self.split = split
self.split_path = data_root / language / split
self.tokenizer = tokenizer
csv_path = self.split_path / f"{split}.csv"
if not csv_path.exists():
raise FileNotFoundError(f"Missing CSV file: {csv_path}")
df = pd.read_csv(csv_path).dropna(axis=0).copy()
required_cols = {image_id_col, text_col, label_col}
missing_cols = required_cols.difference(df.columns)
if missing_cols:
raise ValueError(f"{csv_path} is missing columns: {sorted(missing_cols)}")
df[image_id_col] = df[image_id_col].astype(int)
df[label_col] = df[label_col].map(normalize_label)
id_to_row = {int(row[image_id_col]): row for _, row in df.iterrows()}
raw_images: List[torch.Tensor] = []
raw_records: List[Tuple[int, str, float]] = []
for image_file in sorted(self.split_path.iterdir()):
if image_file.suffix.lower() == ".csv":
continue
if image_file.suffix.lower() not in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}:
continue
try:
image_id = int(image_file.stem)
except ValueError:
continue
if image_id not in id_to_row:
continue
image = cv2.imread(str(image_file))
if image is None:
raise ValueError(f"Could not read image: {image_file}")
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = cv2.resize(image, (224, 224), interpolation=cv2.INTER_CUBIC)
image_tensor = torch.tensor(image, dtype=torch.float32).permute(2, 0, 1) / 255.0
raw_images.append(image_tensor)
row = id_to_row[image_id]
raw_records.append((image_id, str(row[text_col]), float(row[label_col])))
if not raw_records:
raise ValueError(f"No valid examples found in {self.split_path}")
image_stack = torch.stack(raw_images, dim=0)
mean = image_stack.mean(dim=(0, 2, 3), keepdim=True)
std = image_stack.std(dim=(0, 2, 3), keepdim=True).clamp_min(1e-6)
image_stack = (image_stack - mean) / std
self.images: List[torch.Tensor] = []
self.labels: List[float] = []
self.image_ids: List[int] = []
texts: List[str] = []
for idx, (image_id, text, label) in enumerate(raw_records):
repeat = 1
if split == "train" and int(label) == 1 and oversample_positive_train > 0:
repeat += oversample_positive_train
for _ in range(repeat):
self.images.append(image_stack[idx])
self.labels.append(label)
self.image_ids.append(image_id)
texts.append(text)
tokens = tokenizer(
texts,
truncation=True,
max_length=max_length,
padding="max_length",
return_tensors="pt",
)
self.input_ids = tokens.input_ids
self.attention_mask = tokens.attention_mask
def __len__(self) -> int:
return len(self.labels)
def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:
return {
"image": self.images[index],
"label": torch.tensor(self.labels[index], dtype=torch.float32),
"input_ids": self.input_ids[index],
"attention_mask": self.attention_mask[index],
"image_id": torch.tensor(self.image_ids[index], dtype=torch.long),
}
|