File size: 5,926 Bytes
fa1a600
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# Copyright 2024 Big Vision Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# pylint: disable=line-too-long
r"""Generates COCO-35L in a TFDS-ready structure.



First, download the captions from https://google.github.io/crossmodal-3600/ and the images from https://cocodataset.org/#download.

The coco Karpathy split is available at http://cs.stanford.edu/people/karpathy/deepimagesent/caption_datasets.zip:

  mkdir -p /tmp/data/coco35l/images

  wget https://storage.googleapis.com/crossmodal-3600/coco_mt_train.jsonl.bz2 -P /tmp/data/coco35l

  wget https://storage.googleapis.com/crossmodal-3600/coco_mt_dev.jsonl.bz2 -P /tmp/data/coco35l

  bzip2 -dk  /tmp/data/coco35l/coco_mt_train.jsonl.bz2 /tmp/data/coco35l/coco_mt_dev.jsonl.bz2

  wget http://cs.stanford.edu/people/karpathy/deepimagesent/caption_datasets.zip -P /tmp/data/coco35l

  unzip /tmp/data/coco35l/caption_datasets.zip -d /tmp/data/coco35l/

  wget http://images.cocodataset.org/zips/train2014.zip -P /tmp/data/coco35l/images

  wget http://images.cocodataset.org/zips/val2014.zip -P /tmp/data/coco35l/images

  unzip /tmp/data/coco35l/images/train2014.zip -d /tmp/data/coco35l/images/

  unzip /tmp/data/coco35l/images/val2014.zip -d /tmp/data/coco35l/images/



Then, run conversion locally (make sure to install tensorflow-datasets for the `tfds` util):



  cd big_vision/datasets

  env TFDS_DATA_DIR=/tmp/tfds tfds build --datasets=coco35l



Example to load:



  import tensorflow_datasets as tfds

  dataset = tfds.load(

      'coco35l', split='dev_en',

      data_dir='/tmp/tfds')

"""

import json
import os.path

import tensorflow_datasets as tfds

_DESCRIPTION = """

COCO image + captions, translated from English to 35 languages (English incl.).

"""

# pylint: disable=line-too-long
_CITATION = """

@inproceedings{thapliyal-etal-2022-crossmodal,

    title = "Crossmodal-3600: A Massively Multilingual Multimodal Evaluation Dataset",

    author = "Thapliyal, Ashish V.  and

      Pont Tuset, Jordi  and

      Chen, Xi  and

      Soricut, Radu",

    editor = "Goldberg, Yoav  and

      Kozareva, Zornitsa  and

      Zhang, Yue",

    booktitle = "Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing",

    month = dec,

    year = "2022",

    address = "Abu Dhabi, United Arab Emirates",

    publisher = "Association for Computational Linguistics",

    url = "https://aclanthology.org/2022.emnlp-main.45",

    doi = "10.18653/v1/2022.emnlp-main.45",

    pages = "715--729",

}

"""
# pylint: enable=line-too-long


_CAPTIONS_PATH = '/tmp/data/coco35l'
_IMAGES_PATH = '/tmp/data/mscoco/images'
_COCOCAPS_PATH = '/tmp/data/mscoco/dataset_coco.json'

LANGUAGES = [
    'ar', 'bn', 'cs', 'da', 'de', 'el', 'en', 'es', 'fa', 'fi', 'fil', 'fr',
    'he', 'hi', 'hr', 'hu', 'id', 'it', 'ja', 'ko', 'mi', 'nl', 'no', 'pl',
    'pt', 'ro', 'ru', 'sv', 'sw', 'te', 'th', 'tr', 'uk', 'vi', 'zh',
]


class Coco35l(tfds.core.GeneratorBasedBuilder):
  """DatasetBuilder for COCO-35L dataset."""

  VERSION = tfds.core.Version('1.0.0')
  RELEASE_NOTES = {'1.0.0': 'First release.'}

  def _info(self):
    """Returns the metadata."""

    return tfds.core.DatasetInfo(
        builder=self,
        description=_DESCRIPTION,
        features=tfds.features.FeaturesDict({
            'image/id': tfds.features.Text(),
            'image': tfds.features.Image(encoding_format='jpeg'),
            'captions': tfds.features.Sequence(tfds.features.Text()),
            'language': tfds.features.Text(),
        }),
        supervised_keys=None,
        homepage='https://google.github.io/crossmodal-3600/',
        citation=_CITATION,
    )

  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
    """Returns SplitGenerators."""
    splits = []
    for lang in LANGUAGES:
      splits.extend([f'train_{lang}', f'dev_{lang}'])
    return {split: self._generate_examples(split) for split in splits}

  def _generate_examples(self, split: str):
    """Yields (key, example) tuples from dataset."""
    split, language = split.split('_')

    id_to_path = dict()
    with open(_COCOCAPS_PATH, 'r') as f:
      data = json.load(f)['images']
    for d in data:
      id_to_path[d['cocoid']] = os.path.join(
          _IMAGES_PATH, d['filepath'], d['filename']
      )

    annot_fname = os.path.join(_CAPTIONS_PATH, f'coco_mt_{split}.jsonl')
    data = {}
    with open(annot_fname, 'r') as f:
      for line in f:
        j = json.loads(line)
        image_id = f'{j["image_id"].split("_")[0]}_{language}'
        if image_id not in data:
          data[image_id] = []
        if language == 'en':
          # COCO-35L was constructed from English into 35 other languages.
          # To add English in our TFDS, we just select a language (eg. "de") to
          # have each unique example, and add the corresponding source caption.
          if j['trg_lang'] == 'de':
            data[image_id].append(j['caption_tokenized'])
        else:
          if j['trg_lang'] == language:
            data[image_id].append(j['translation_tokenized'])

    for image_id, captions in data.items():
      yield image_id, {
          'image/id': image_id,
          'image': id_to_path[int(image_id.split('_')[0])],
          'captions': captions,
          'language': language,
      }