Spaces:
Running on Zero
Running on Zero
File size: 2,457 Bytes
a460c95 | 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 | from pathlib import Path
import json
from PIL import Image
import torch
from torch.utils.data import Dataset
# COCO2014 validation set
class COCOValDataset(Dataset):
def __init__(self, path, transform):
super().__init__()
self.path = Path(path)
self.image_files = list(self.path.iterdir())
self.transform = transform
def __len__(self):
return len(self.image_files)
def __getitem__(self, index):
with open(self.image_files[index], "rb") as f:
image = Image.open(f)
image = image.convert("RGB")
return {"image": self.transform(image), "image_id": self.image_files[index].name[:-4]}
@staticmethod
def collate_fn(batch):
batch_size = len(batch)
images = []
image_ids = []
for i in range(batch_size):
images.append(batch[i]["image"])
image_ids.append(batch[i]["image_id"])
return {
"images": torch.stack(images, dim=0),
"image_ids": image_ids,
}
# samples randomly sampled from COCO2014 validation set
# COCO30K or COCO6K
class COCODataset(Dataset):
def __init__(self, path, debug=False):
super().__init__()
self.path = Path(path)
self.captions = json.load(open(self.path, "r"))["annotations"]
if debug:
self.captions = self.captions[:1024]
def __len__(self):
# return 120 # for test
return len(self.captions)
def __getitem__(self, index):
return {
"id": self.captions[index]["id"],
# different captions may correspond to the same image_id
"image_id": self.captions[index]["image_id"],
"type": "prompt",
"input": self.captions[index]["caption"],
"seed": self.captions[index]["seed"],
}
@staticmethod
def collate_fn(batch):
batch_size = len(batch)
ids = []
image_ids = []
types = []
inputs = []
seeds = []
for i in range(batch_size):
ids.append(batch[i]["id"])
image_ids.append(batch[i]["image_id"])
types.append(batch[i]["type"])
inputs.append(batch[i]["input"])
seeds.append(batch[i]["seed"])
return {
"ids": ids,
"image_ids": image_ids,
"type": types,
"input": inputs,
"seeds": seeds,
}
|