diffTell_dataloader.py

import json
from pathlib import Path
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms


class ImagePairConversationDataset(Dataset):
    def __init__(self, json_path, transform=None):
        self.data = json.loads(Path(json_path).read_text())
        self.transform = transform or transforms.Compose(
            [transforms.Resize((224, 224)), transforms.ToTensor()]
        )

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

    def __getitem__(self, idx):
        item = self.data[idx]

        # img1 = Image.open(item["images"][0]).convert("RGB")
        # img2 = Image.open(item["images"][1]).convert("RGB")

        # img1 = self.transform(img1)
        # img2 = self.transform(img2)
        img1 = item["images"][0]
        img2 = item["images"][1]

        user_query = next(
            msg["content"] for msg in item["conversation"] if msg["role"] == "user"
        )
        assistant_response = next(
            msg["content"] for msg in item["conversation"] if msg["role"] == "assistant"
        )

        return {
            "id": item["id"],
            "image1": img1,
            "image2": img2,
            "question": user_query,
            "answer": assistant_response,
        }


# Example usage
if __name__ == "__main__":
    dataset = ImagePairConversationDataset("diffTell_train_with_source.json")
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)

    for batch in dataloader:
        print(batch["id"])
        print(batch["image1"], batch["image2"])
        print(batch["question"])
        print(batch["answer"])
        break  # Just show the first batch