Spaces:
Runtime error
Runtime error
File size: 2,331 Bytes
03561be |
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 |
import json
import os
import random
from collections import defaultdict
from PIL import Image
from .vqa_dataset import VQADataset
class CLEVRDataset(VQADataset):
"""Visual Reasoning Dataset. It also contains Dialog.
Note: The image is a little bit simple. with several objects and simple background.
"""
def __init__(self, tokenizer, vis_processor, vis_root, ann_paths, **kwargs):
super().__init__(tokenizer, vis_processor, vis_root, ann_paths=[], **kwargs)
self.annotation = self.load_annotations(ann_paths)
if self.sample_image:
print("randomly sample one annotation for each image")
self.annotation = self.parse_annotation(self.annotation)
self._add_instance_ids()
@staticmethod
def load_annotations(ann_paths):
annotation = []
for ann_path in ann_paths:
ann = json.load(open(ann_path, "r"))
annotation.extend(ann["questions"])
return annotation
def parse_annotation(self, annotation):
image_list = defaultdict(list)
for ann in annotation:
image_list[ann["image_filename"]].append(ann)
annotation = []
for ann_list in image_list.values():
annotation.append(random.choice(ann_list))
return annotation
def process_text(self, ann):
question = ann["question"]
answer = ann["answer"]
instruction = self.prompter(question)
return dict(instruction=instruction, answer=answer)
def process_image(self, ann):
split = ann["split"]
image_path = os.path.join(self.vis_root, split, ann["image_filename"])
image = Image.open(image_path).convert("RGB")
image = self.vis_processor(image)
return image
def build_clevr_dataset(
tokenizer,
vis_processor,
vis_root="data/clevr/CLEVR_v1.0/images",
ann_paths=[
"data/clevr/CLEVR_v1.0/questions/CLEVR_train_questions.json",
"data/clevr/CLEVR_v1.0/questions/CLEVR_val_questions.json",
],
sample_image=False,
):
return CLEVRDataset(
tokenizer=tokenizer,
vis_processor=vis_processor,
vis_root=vis_root,
ann_paths=ann_paths,
sample_image=sample_image,
)
|