File size: 7,280 Bytes
9628bd8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
import csv
import json
import os

import datasets
import pandas as pd
import numpy as np


class ImageCaptionBuilderConfig(datasets.BuilderConfig):

    def __init__(self, name, splits, langs, prefix_before_image_fn=False, zfill=1, **kwargs):

        super().__init__(name, **kwargs)

        self.splits = splits
        self.langs = langs
        self.prefix_before_image_fn = prefix_before_image_fn
        self.zfill = zfill


# TODO: Add BibTeX citation
# Find for instance the citation on arxiv or on the dataset repo/website
_CITATION = """\
@InProceedings{None,
    title = {Generic images to captions dataset},
    author={Yih-Dar SHIEH},
    year={2020}
}
"""

# TODO: Add description of the dataset here
# You can copy an official description
_DESCRIPTION = """\

"""

# TODO: Add a link to an official homepage for the dataset here
_HOMEPAGE = ""

# TODO: Add the licence for the dataset here if you can find it
_LICENSE = ""

# TODO: Add link to the official dataset URLs here
# The HuggingFace dataset library don't host the datasets but only point to the original files
# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)
_URLs = {}


# TODO: Name of the dataset usually match the script name with CamelCase instead of snake_case
class ImageCaptionDataset(datasets.GeneratorBasedBuilder):
    """TODO: Short description of my dataset."""

    VERSION = datasets.Version("0.0.0")

    BUILDER_CONFIG_CLASS = ImageCaptionBuilderConfig
    BUILDER_CONFIGS = [
        ImageCaptionBuilderConfig(name='coco_2017', splits=['train', 'valid'], prefix_before_image_fn=False, zfill=12, langs=['en', 'fr']),
        ImageCaptionBuilderConfig(name='cc3m', splits=['train', 'valid'], prefix_before_image_fn=True, zfill=8, langs=['en', 'fr']),
        ImageCaptionBuilderConfig(name='cc12m', splits=['train', 'valid'], prefix_before_image_fn=True, zfill=8, langs=['en', 'fr'])
    ]
    DEFAULT_CONFIG_NAME = "coco_2017"

    def _info(self):
        # TODO: This method specifies the datasets.DatasetInfo object which contains informations and typings for the dataset

        feature_dict = {
            "image_id": datasets.Value("int64"),
            "id": datasets.Value("int64"),
            "caption": datasets.Value("string"),
        }
        for lang in self.config.langs:
            feature_dict[lang] = datasets.Value("string")
        feature_dict["image_url"] = datasets.Value("string")
        feature_dict["image_file"] = datasets.Value("string")

        features = datasets.Features(feature_dict)

        return datasets.DatasetInfo(
            # This is the description that will appear on the datasets page.
            description=_DESCRIPTION,
            # This defines the different columns of the dataset and their types
            features=features,  # Here we define them above because they are different between the two configurations
            # If there's a common (input, target) tuple from the features,
            # specify them here. They'll be used if as_supervised=True in
            # builder.as_dataset.
            supervised_keys=None,
            # Homepage of the dataset for documentation
            homepage=_HOMEPAGE,
            # License for the dataset if available
            license=_LICENSE,
            # Citation for the dataset
            citation=_CITATION,
        )

    def _split_generators(self, dl_manager):
        """Returns SplitGenerators."""
        # TODO: This method is tasked with downloading/extracting the data and defining the splits depending on the configuration
        # If several configurations are possible (listed in BUILDER_CONFIGS), the configuration selected by the user is in self.config.name

        data_dir = self.config.data_dir

        splits = []
        for split in self.config.splits:
            if split == 'train':
                dataset = datasets.SplitGenerator(
                    name=datasets.Split.TRAIN,
                    # These kwargs will be passed to _generate_examples
                    gen_kwargs={
                        "jsonl_dir": os.path.join(data_dir, f'{self.config.name}_jsonls', 'train'),
                        "image_dir": os.path.join(data_dir, f'{self.config.name}_images', 'train'),
                        "split": "train",
                    }
                )
            elif split in ['val', 'valid', 'validation', 'dev']:
                dataset = datasets.SplitGenerator(
                    name=datasets.Split.VALIDATION,
                    # These kwargs will be passed to _generate_examples
                    gen_kwargs={
                        "jsonl_dir": os.path.join(data_dir, f'{self.config.name}_jsonls', 'valid'),
                        "image_dir": os.path.join(data_dir, f'{self.config.name}_images', 'valid'),
                        "split": "valid",
                    },
                )
            elif split == 'test':
                dataset = datasets.SplitGenerator(
                    name=datasets.Split.TEST,
                    # These kwargs will be passed to _generate_examples
                    gen_kwargs={
                        "jsonl_dir": os.path.join(data_dir, f'{self.config.name}_jsonls', 'test'),
                        "image_dir": os.path.join(data_dir, f'{self.config.name}_images', 'test'),
                        "split": "test",
                    },
                )
            else:
                continue

            splits.append(dataset)

        return splits

    def _generate_examples(
        # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`
        self, jsonl_dir, image_dir, split
    ):
        """ Yields examples as (key, example) tuples. """
        # This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.
        # The `key` is here for legacy reason (tfds) and is not important in itself.

        if split == 'dev':
            split = 'valid'

        fns = [os.path.join(jsonl_dir, fn) for fn in os.listdir(jsonl_dir) if os.path.isfile(os.path.join(jsonl_dir, fn)) and fn.endswith("jsonl")]

        for jsonl_file in fns:

            with open(jsonl_file, 'r', encoding='UTF-8') as fp:

                for id_, line in enumerate(fp):

                    ex = json.loads(line)

                    example = {
                        "image_id": ex['image_id'],
                        "id": ex["id"],
                        "caption": ex["caption"],
                    }

                    for lang in self.config.langs:
                        example[lang] = ex[lang]

                    if 'image_url' in ex:
                        example['image_url'] = ex['image_url']
                    else:
                        example['image_url'] = ''

                    fn = f'{str(ex["image_id"]).zfill(self.config.zfill)}.jpg'
                    if self.config.prefix_before_image_fn:
                        fn = f'{self.config.name}_{split}_' + fn

                    image_file = os.path.join(image_dir, fn)
                    example['image_file'] = image_file

                    if not os.path.isfile(image_file):
                        continue

                    yield id_, example