File size: 4,210 Bytes
b91f38a | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | import os
import torch
import json
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms
class SimpleImageTextDataset(Dataset):
def __init__(
self,
data_source, # ✅ 可以是 CSV / list / all.jsonl
tokenizer,
image_size=224,
use_modal_embedding=False,
default_modal_index=0,
dataset_root="" # ✅ 给 all.jsonl 用
):
self.tokenizer = tokenizer
self.use_modal_embedding = use_modal_embedding
self.default_modal_index = default_modal_index
self.image_size = image_size
self.dataset_root = dataset_root
# ------- 读取样本 -------
if isinstance(data_source, str):
# ✅ 1️⃣ 如果是 CSV
if data_source.endswith(".csv"):
import pandas as pd
df = pd.read_csv(data_source)
self.samples = list(zip(df["image_path"], df["text"]))
# ✅ 2️⃣ 如果是 all.jsonl
elif data_source.endswith(".jsonl"):
self.samples = self._load_from_jsonl(data_source)
else:
raise ValueError("Unsupported file type. Use CSV or JSONL.")
else:
# ✅ 3️⃣ 直接是 list [(img_path, text), ...]
self.samples = data_source
# ------- 图像 transform -------
self.transform = transforms.Compose([
transforms.Grayscale(num_output_channels=1),
transforms.Resize((image_size, image_size)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5], std=[0.5])
])
print(f"✅ Loaded {len(self.samples)} image-text pairs.")
def _load_from_jsonl(self, jsonl_path):
samples = []
with open(jsonl_path, "r", encoding="utf-8") as f:
for line in f:
item = json.loads(line.strip())
rel_path = item["img_path"]
rel_path = rel_path.replace("\\", "/")
if rel_path.startswith("dataset/"):
rel_path = rel_path[len("dataset/"):]
if self.dataset_root:
img_path = os.path.join(self.dataset_root, rel_path)
else:
img_path = rel_path
text = item["text"].replace("\n", " ").strip()
samples.append((img_path, text))
return samples
def __len__(self):
return len(self.samples)
def __getitem__(self, index):
img_path, text_str = self.samples[index]
# ---- 1. 载入图像 -> [1,224,224] ----
img = Image.open(img_path).convert("L")
img = self.transform(img)
# ---- 2. tokenize 文本 ----
text = self.tokenizer(
text_str,
return_tensors="pt",
truncation=True,
padding="max_length",
max_length=512,
)
text = {k: v.squeeze(0) for k, v in text.items()}
# ---- 3. modal_embedding（可选）----
if self.use_modal_embedding:
modal_index = torch.tensor(self.default_modal_index, dtype=torch.long)
return {
"image": img,
"text": text,
"modal_index": modal_index
}
return {
"image": img,
"text": text
}
class TextBatch:
def __init__(self, d):
for k, v in d.items():
setattr(self, k, v)
def simple_collate(batch):
images = torch.stack([item["image"] for item in batch], dim=0)
text_dict = {
"input_ids": torch.stack([item["text"]["input_ids"] for item in batch]),
"attention_mask": torch.stack([item["text"]["attention_mask"] for item in batch]),
}
if "token_type_ids" in batch[0]["text"]:
text_dict["token_type_ids"] = torch.stack(
[item["text"]["token_type_ids"] for item in batch]
)
text = TextBatch(text_dict)
if "modal_index" in batch[0]:
modal_indexs = torch.stack([item["modal_index"] for item in batch])
return images, text, modal_indexs
return images, text
|