File size: 1,603 Bytes
3ea5987 | 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 | import os
import pandas as pd
from PIL import Image
import torch
from torch.utils.data import (
DataLoader,
Dataset,
IterableDataset,
SubsetRandomSampler,
get_worker_info,
)
import clip.clip as clip
class CsvDataset(Dataset):
def __init__(self, input_filename, transforms, img_key, caption_key, sep="\t"):
df = pd.read_csv(input_filename, sep=sep)
self.location = os.path.dirname(input_filename)
self.images = df[img_key].tolist()
self.captions = df[caption_key].tolist()
self.transforms = transforms
def __len__(self):
return len(self.captions)
def __getitem__(self, idx):
image_path = os.path.join(self.location, str(self.images[idx]))
images = self.transforms(Image.open(image_path))
texts = clip.tokenize([str(self.captions[idx])])[0]
return images, texts
class conceptual_captions(Dataset):
def __init__(
self, transforms, location, batch_size, *args, num_workers=16, **kwargs
):
file_name = "Validation_GCC-1.1.0-Validation_output.csv"
file_path = os.path.join(location, file_name)
self.template = lambda c: f"a photo of a {c}."
self.train_dataset = CsvDataset(
input_filename=file_path,
transforms=transforms,
img_key="filepath",
caption_key="title",
)
# breakpoint()
self.train_loader = torch.utils.data.DataLoader(
self.train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
)
|