Add files using upload-large-folder tool

Tipsomaly/model/big_vision/datasets/ai2d/ai2d.py

@@ -0,0 +1,209 @@
+# 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"""AI2D TFDS converter.
+
+
+It's a small dataset, so can be built locally. Copy the data to local disk:
+
+  mkdir -p /tmp/data/ai2d && cd /tmp/data/ai2d
+  wget https://ai2-public-datasets.s3.amazonaws.com/diagrams/ai2d-all.zip
+  wget https://s3-us-east-2.amazonaws.com/prior-datasets/ai2d_test_ids.csv
+  wget https://github.com/googlefonts/dm-fonts/raw/main/Sans/fonts/ttf/DMSans-Regular.ttf
+  unzip ai2d-all.zip
+
+Also download a font for rendering, set the location in the flag font_path.
+
+Then, run conversion locally (make sure to install tensorflow-datasets for the `tfds` util):
+
+  cd third_party/py/big_vision/datasets
+  env TFDS_DATA_DIR=/tmp/tfds tfds build --datasets=ai2d
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset = tfds.load(ai2d', split='train', data_dir='/tmp/tfds')
+"""
+
+import functools
+import glob
+import io
+import json
+import os
+from typing import Any, Dict
+
+from absl import flags
+import numpy as np
+from PIL import Image
+from PIL import ImageDraw
+from PIL import ImageFont
+import tensorflow_datasets as tfds
+
+
+_DESCRIPTION = """AI2D dataset."""
+
+# pylint: disable=line-too-long
+_CITATION = """
+@inproceedings{kembhavi2016eccv,
+  author = {Aniruddha Kembhavi, Mike Salvato, Eric Kolve, Minjoon Seo, Hannaneh Hajishirzi, Ali Farhadi},
+  title = {A Diagram Is Worth A Dozen Images},
+  booktitle = {European Conference on Computer Vision (ECCV)},
+  year = {2016}
+  url={https://api.semanticscholar.org/CorpusID:2682274}
+}
+"""
+# pylint: enable=line-too-long
+
+
+_INPUT_PATH = flags.DEFINE_string(
+    'input_path', '/tmp/data/ai2d/', 'Downloaded AI2D data.'
+)
+_FONT_PATH = flags.DEFINE_string(
+    'font_path', '/tmp/data/ai2d/DMSans-Regular.ttf',
+    'Font for rendering annotations.'
+)
+
+
+class Ai2d(tfds.core.GeneratorBasedBuilder):
+  """DatasetBuilder for AI2D dataset."""
+
+  VERSION = tfds.core.Version('1.1.0')
+  RELEASE_NOTES = {'1.1.0': 'Re-create from scratch + more fields.'}
+
+  def _info(self):
+    """Returns the metadata."""
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=_DESCRIPTION,
+        features=tfds.features.FeaturesDict({
+            'id': tfds.features.Text(),
+            'question': tfds.features.Text(),
+            'label': tfds.features.Scalar(np.int32),
+            'answer': tfds.features.Text(),
+            'possible_answers': tfds.features.Sequence(tfds.features.Text()),
+            'abc_label': tfds.features.Scalar(np.bool_),
+            'image_name': tfds.features.Text(),
+            'image': tfds.features.Image(encoding_format='png'),
+        }),
+        homepage='https://allenai.org/data/diagrams',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    return {split: self._generate_examples(split)
+            for split in ('test', 'train')}
+
+  def _generate_examples(self, split: str):
+    """Yields (key, example) tuples."""
+    with open(
+        os.path.join(_INPUT_PATH.value, 'ai2d_test_ids.csv'), 'r'
+    ) as f:
+      all_test_ids = f.readlines()
+    all_test_ids = [line.strip() for line in all_test_ids]
+
+    all_annotation_paths = glob.glob(
+        os.path.join(_INPUT_PATH.value, 'ai2d/questions', '*.json'))
+    for annotation_path in all_annotation_paths:
+      basename = os.path.basename(annotation_path)
+      image_id = basename.split('.')[0]
+      if image_id in all_test_ids and split == 'train':
+        continue
+      elif image_id not in all_test_ids and split == 'test':
+        continue
+
+      text_annotation_path = os.path.join(
+          _INPUT_PATH.value, 'ai2d/annotations', basename
+      )
+      with open(annotation_path, 'r') as f:
+        with open(text_annotation_path, 'r') as g:
+          question_json = json.load(f)
+          text_annotation_json = json.load(g)
+          for question in question_json['questions']:
+            label_id = int(
+                question_json['questions'][question]['correctAnswer']
+            )
+            choices = question_json['questions'][question]['answerTexts']
+            abc_label = question_json['questions'][question]['abcLabel']
+            annotation = {
+                'id': question_json['questions'][question]['questionId'],
+                'question': question,
+                'label': label_id,
+                'answer': choices[label_id],
+                'possible_answers': tuple(choices),
+                'abc_label': abc_label,
+                'image_name': question_json['imageName'],
+            }
+            annotation['image'] = _create_image(
+                annotation, text_annotation_json['text']
+            )
+            yield annotation['id'], annotation
+
+
+@functools.cache
+def Font(  # pylint: disable=invalid-name
+    size: int,
+) -> ImageFont.FreeTypeFont:
+  """Loads the font from in the specified style.
+
+  Args:
+    size: The size of the returned font.
+
+  Returns:
+    The loaded font.
+  """
+  return ImageFont.truetype(_FONT_PATH.value, size=size)
+
+
+def _create_image(
+    annotation: Dict[str, Any], text_annotation: Dict[str, Any]
+) -> bytes:
+  """Adds image to one annotation."""
+  img_path = os.path.join(_INPUT_PATH.value, 'ai2d/images',
+                          annotation['image_name'])
+  with open(img_path, 'rb') as f:
+    if annotation['abc_label']:
+      raw_image = _draw_text(f, text_annotation)
+    else:
+      raw_image = f.read()
+  return raw_image
+
+
+def _draw_text(image, text_annotations) -> bytes:
+  """Replaces text in image by the correct replacement letter from AI2D."""
+  image = Image.open(image)
+  draw = ImageDraw.Draw(image)
+  for annotation in text_annotations:
+    current_annotation = text_annotations[annotation]
+    rectangle = current_annotation['rectangle']
+    box = [tuple(rectangle[0]), tuple(rectangle[1]),]
+    text = current_annotation['replacementText']
+    position = box[0]
+    draw.rectangle(box, fill='white')
+    font_size = 100
+    x_diff = box[1][0] - box[0][0]
+    y_diff = box[1][1] - box[0][1]
+    font = Font(font_size)
+    size = font.getbbox(text)
+    while (size[2] > x_diff or size[3] > y_diff) and font_size > 0:
+      font = Font(font_size)
+      size = font.getbbox(text)
+      font_size -= 1
+    delta = (x_diff - size[2]) // 2
+    position = (position[0] + delta, position[1])
+    draw.text(position, text, fill='black', font=font)
+  new_image_bytes = io.BytesIO()
+  image.save(new_image_bytes, format='PNG')
+  return new_image_bytes.getvalue()

Tipsomaly/model/big_vision/datasets/aokvqa/aokvqa.py

@@ -0,0 +1,182 @@
+# 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"""Implements the OKVQA dataset for TFDS.
+
+Download the required files from https://aokvqa.allenai.org/download.html:
+
+mkdir -p /tmp/tfds
+cd /tmp/tfds/
+wget http://images.cocodataset.org/zips/train2017.zip
+wget http://images.cocodataset.org/zips/val2017.zip
+wget http://images.cocodataset.org/zips/test2017.zip
+wget https://prior-datasets.s3.us-east-2.amazonaws.com/aokvqa/aokvqa_v1p0.tar.gz
+unzip val2017.zip
+unzip train2017.zip
+unzip test2017.zip
+tar xzf aokvqa_v1p0.tar.gz
+
+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=aokvqa
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset = tfds.load('aokvqa', split='val', data_dir='/tmp/tfds')
+"""
+
+import json
+import os
+from typing import Any
+import numpy as np
+import tensorflow_datasets as tfds
+
+_DESCRIPTION = """
+A-OKVQA addresses the task of VQA with outside knowledge.
+It is a follow-up dataset of OKVQA.
+
+This version of the dataset contains:
+- Questions + Answers + Multiple Choice Answers + Rationales from A-OKVQA.
+- Images from COCO.
+"""
+
+_CITATION = """
+@article{AOKVQA,
+  title={A-OKVQA: A Benchmark for Visual Question Answering using World Knowledge},
+  author={Dustin Schwenk and Apoorv Khandelwal and Christopher Clark and Kenneth Marino and Roozbeh Mottaghi},
+  journal={arXiv},
+  year={2022},
+}
+"""
+
+ANNOTATION_FILES = {
+    'train': 'aokvqa_v1p0_train.json',
+    'val': 'aokvqa_v1p0_val.json',
+    'test': 'aokvqa_v1p0_test.json',
+}
+
+
+# When running locally (recommended), copy files as above an use these:
+_AOKVQA_PATH = '/tmp/tfds'
+
+
+class AOkVqa(tfds.core.GeneratorBasedBuilder):
+  """AOKVQA dataset for TFDS."""
+
+  VERSION = tfds.core.Version('1.0.0')
+  RELEASE_NOTES = {'1.0.0': 'ArrayRecord version.'}
+  MANUAL_DOWNLOAD_INSTRUCTIONS = """
+  In manual_dir/ you should have a directory a_ok_vqa which contains the
+  following files and directories:
+  From the A-OKVQA dataset:
+  - aokvqa_v1p0_train.json
+  - aokvqa_v1p0_val.json
+  - aokvqa_v1p0_test.json
+  It also requires the COCO data files.
+  """
+
+  def _info(self) -> tfds.core.DatasetInfo:
+    """Returns the dataset metadata."""
+    features = tfds.features.FeaturesDict({
+        'image': tfds.features.Image(shape=(None, None, 3)),
+        'image_id': tfds.features.Scalar(dtype=np.int64),
+        'direct_answers': tfds.features.Sequence(tfds.features.Text()),
+        'direct_answer_is_difficult': tfds.features.Scalar(dtype=np.bool_),
+        'multiple_choice_possible_answers':  # List of 4 possible answers.
+            tfds.features.Sequence(tfds.features.Text()),
+        'multiple_choice_correct_idx':  # Integer from 0-3.
+            tfds.features.Scalar(dtype=np.int32),
+        'answer_rationales': tfds.features.Sequence(tfds.features.Text()),
+        'question': tfds.features.Text(),
+        'question_id': tfds.features.Text(),
+    })
+
+    return tfds.core.DatasetInfo(
+        builder=self,
+        features=features,
+        description=_DESCRIPTION,
+        supervised_keys=None,
+        homepage='https://okvqa.allenai.org/',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager) -> ...:
+    """Call the function which defines the splits."""
+    # data_dir = dl_manager.manual_dir
+    data_dir = _AOKVQA_PATH
+    return {
+        'train': self._generate_examples(data_dir, 'train'),
+        'val': self._generate_examples(data_dir, 'val'),
+        'test': self._generate_examples(data_dir, 'test'),
+    }
+
+  def _generate_examples(self, data_dir: str, split: str) -> ...:
+    annotations = get_annotations(data_dir, split)
+
+    for question_id, feature_dict in annotations.items():
+      image_id = feature_dict['image_id']
+
+      # Add image and GT segmentatio labels from total_transfer.
+      feature_dict['image'] = self.get_image_path(data_dir, split, image_id)
+
+      # Add dummy features for several features in the test set.
+      if split not in ['train', 'val']:
+        assert split == 'test', f'Unknown split: {split}'
+        feature_dict['multiple_choice_correct_idx'] = -1
+        feature_dict['direct_answers'] = []
+        feature_dict['answer_rationales'] = []
+      yield f'{question_id}', feature_dict
+
+  def get_image_path(self, data_dir: str, split: str, image_id: int) -> str:
+    return f'{data_dir}/{split}2017/{image_id:012d}.jpg'
+
+
+def get_annotations(
+    data_dir: str, split: str) -> dict[int, dict[str, Any]]:
+  """Return okvqa annotations (quesions and answers) as dictionary."""
+  path = os.path.join(data_dir, ANNOTATION_FILES[split])
+  with open(path) as f:
+    annotations = json.load(f)
+
+  aokvqa_annotations = {}
+  for annotation in annotations:
+    # Sanity checks
+    assert len(annotation['choices']) == 4
+
+    question_id = annotation['question_id']
+
+    aokvqa_annotations[question_id] = {
+        'image_id': annotation['image_id'],
+        'direct_answer_is_difficult': annotation['difficult_direct_answer'],
+        'multiple_choice_possible_answers': annotation['choices'],
+        'question': annotation['question'],
+        'question_id': annotation['question_id'],
+    }
+
+    # Get answers and rationales for train and val only, not for test.
+    if split in ['train', 'val']:
+      assert len(annotation['direct_answers']) == 10
+      assert len(annotation['rationales']) == 3
+
+      aokvqa_annotations[question_id]['direct_answers'] = annotation[
+          'direct_answers']
+      aokvqa_annotations[question_id]['answer_rationales'] = annotation[
+          'rationales']
+      aokvqa_annotations[question_id]['multiple_choice_correct_idx'] = (
+          annotation['correct_choice_idx'])
+
+  return aokvqa_annotations

Tipsomaly/model/big_vision/datasets/chartqa/chartqa.py

@@ -0,0 +1,122 @@
+# 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"""Implements CharQA in TFDS structure.
+
+It's small data, so simple to run locally. First, copy the data to local disk:
+
+  mkdir -p /tmp/data
+  wget -O /tmp/data/chartqa.zip https://huggingface.co/datasets/ahmed-masry/ChartQA/resolve/main/ChartQA%20Dataset.zip?download=true
+  unzip /tmp/data/chartqa.zip
+
+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=chartqa
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset_augmented = tfds.load('chartqa/augmented', split='train', data_dir='/tmp/tfds')
+"""
+import json
+import os
+
+import numpy as np
+import tensorflow_datasets as tfds
+
+
+_DESCRIPTION = """ChartQA dataset."""
+
+# pylint: disable=line-too-long
+_CITATION = """
+@inproceedings{masry-etal-2022-chartqa,
+    title = "{C}hart{QA}: A Benchmark for Question Answering about Charts with Visual and Logical Reasoning",
+    author = "Masry, Ahmed  and
+      Do, Xuan Long  and
+      Tan, Jia Qing  and
+      Joty, Shafiq  and
+      Hoque, Enamul",
+    editor = "Muresan, Smaranda  and
+      Nakov, Preslav  and
+      Villavicencio, Aline",
+    booktitle = "Findings of the Association for Computational Linguistics: ACL 2022",
+    month = may,
+    year = "2022",
+    address = "Dublin, Ireland",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2022.findings-acl.177",
+    doi = "10.18653/v1/2022.findings-acl.177",
+    pages = "2263--2279",
+    abstract = "Charts are very popular for analyzing data. When exploring charts, people often ask a variety of complex reasoning questions that involve several logical and arithmetic operations. They also commonly refer to visual features of a chart in their questions. However, most existing datasets do not focus on such complex reasoning questions as their questions are template-based and answers come from a fixed-vocabulary. In this work, we present a large-scale benchmark covering 9.6K human-written questions as well as 23.1K questions generated from human-written chart summaries. To address the unique challenges in our benchmark involving visual and logical reasoning over charts, we present two transformer-based models that combine visual features and the data table of the chart in a unified way to answer questions. While our models achieve the state-of-the-art results on the previous datasets as well as on our benchmark, the evaluation also reveals several challenges in answering complex reasoning questions.",
+}
+"""
+# pylint: enable=line-too-long
+
+# When running locally (recommended), copy files as above an use these:
+_CHARTQA_PATH = '/tmp/data/ChartQA Dataset/'
+
+
+class ChartQAConfig(tfds.core.BuilderConfig):
+  """Configuration to build the dataset."""
+  pass
+
+
+class ChartQA(tfds.core.GeneratorBasedBuilder):
+  """DatasetBuilder for ChartQA dataset."""
+
+  VERSION = tfds.core.Version('1.0.0')
+  RELEASE_NOTES = {'1.0.0': 'First release.'}
+  BUILDER_CONFIGS = [
+      ChartQAConfig(name='human', description='Human set'),
+      ChartQAConfig(name='augmented', description='Augmented set'),
+  ]
+
+  def _info(self):
+    """Returns the metadata."""
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=_DESCRIPTION,
+        features=tfds.features.FeaturesDict({
+            'question_id': tfds.features.Scalar(np.int32),
+            'image/filename': tfds.features.Text(),
+            'image': tfds.features.Image(encoding_format='png'),
+            'question': tfds.features.Text(),
+            'answer': tfds.features.Text(),
+        }),
+        homepage='https://github.com/vis-nlp/ChartQA',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    return {split: self._generate_examples(split, self.builder_config.name)
+            for split in ('val', 'train', 'test')}
+
+  def _generate_examples(self, split: str, source: str):
+    """Yields (key, example) tuples from test set."""
+    annot_fname = os.path.join(_CHARTQA_PATH, split, f'{split}_{source}.json')
+
+    with open(annot_fname, 'r') as f:
+      data = json.loads(f.read())
+
+    for idx, v in enumerate(data):
+      yield idx, {
+          'question_id': idx,
+          'image/filename': v['imgname'],
+          'image': os.path.join(_CHARTQA_PATH, split, 'png', v['imgname']),
+          'question': v['query'],
+          'answer': v['label'],
+      }

Tipsomaly/model/big_vision/datasets/coco35l/coco35l.py

@@ -0,0 +1,154 @@
+# 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,
+      }

Tipsomaly/model/big_vision/datasets/countbenchqa/countbenchqa.py

@@ -0,0 +1,164 @@
+# 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
+"""Import CountBenchQA dataset (CountBench dataset with added QA annotations).
+
+It's small data, so simple to run locally. First, download all the data:
+
+  mkdir /tmp/data/ ; cd /tmp/data
+  wget https://huggingface.co/datasets/nielsr/countbench/resolve/main/data/train-00000-of-00001-cf54c241ba947306.parquet
+  wget https://raw.githubusercontent.com/teaching-clip-to-count/teaching-clip-to-count.github.io/main/CountBench.json
+
+Then, update the PATHs below and run conversion locally like so:
+
+  cd big_vision/datasets
+  env TFDS_DATA_DIR=/tmp/tfds tfds build --datasets=countbenchqa
+
+The dataset contains 540 images so the dataset creation is very quick.
+
+There is a single split called huggingface to denote that the images come from
+the hugginface parquet file.
+"""
+
+import io
+import json
+
+import numpy as np
+import pandas as pd
+import PIL
+import tensorflow_datasets as tfds
+
+
+# Huggingface dataset path; this is missing about 10% of the images.
+_COUNTBENCH_PARQUET_PATH = '/tmp/data/train-00000-of-00001-cf54c241ba947306.parquet'
+# Public path to the original CountBench JSON file.
+_COUNTBENCH_JSON_PATH = '/tmp/data/CountBench.json'
+# VQA annotations
+_QA_JSON_PATH = 'countbenchqa/data/countbench_paired_questions.json'
+
+_DESCRIPTION = """
+CountBench: We introduce a new object counting benchmark called CountBench,
+  automatically curated (and manually verified) from the publicly available
+  LAION-400M image-text dataset. CountBench contains a total of 540 images
+  containing between two and ten instances of a particular object, where their
+  corresponding captions reflect this number.
+
+CountBenchQA: Each image is paired with a manually generated question about the
+  number of objects in the image to turn CountBench into a VQA task.
+"""
+
+_CITATION = """
+@article{beyer2024paligemma,
+      title={{PaliGemma: A versatile 3B VLM for transfer}},
+      author={Lucas Beyer and Andreas Steiner and André Susano Pinto and Alexander Kolesnikov and Xiao Wang and Daniel Salz and Maxim Neumann and Ibrahim Alabdulmohsin and Michael Tschannen and Emanuele Bugliarello and Thomas Unterthiner and Daniel Keysers and Skanda Koppula and Fangyu Liu and Adam Grycner and Alexey Gritsenko and Neil Houlsby and Manoj Kumar and Keran Rong and Julian Eisenschlos and Rishabh Kabra and Matthias Bauer and Matko Bošnjak and Xi Chen and Matthias Minderer and Paul Voigtlaender and Ioana Bica and Ivana Balazevic and Joan Puigcerver and Pinelopi Papalampidi and Olivier Henaff and Xi Xiong and Radu Soricut and Jeremiah Harmsen and Xiaohua Zhai},
+      year={2024},
+      journal={arXiv preprint arXiv:2407.07726}
+}
+
+@article{paiss2023countclip,
+      title={{Teaching CLIP to Count to Ten}},
+      author={Paiss, Roni and Ephrat, Ariel and Tov, Omer and Zada, Shiran and Mosseri, Inbar and Irani, Michal and Dekel, Tali},
+      year={2023},
+      journal={arXiv preprint arXiv:2302.12066}
+}
+"""
+
+_HOMEPAGE = 'https://teaching-clip-to-count.github.io/'
+
+
+class CountbenchQA(tfds.core.GeneratorBasedBuilder):
+  """Create CountbenchQA dataset."""
+
+  VERSION = tfds.core.Version('1.2.0')
+  RELEASE_NOTES = {'1.1.0': 'Add `huggingface` split.',
+                   '1.2.0': 'Fix image loading for `huggingface` split.'}
+  MANUAL_DOWNLOAD_INSTRUCTIONS = """
+  There are two parts which should be downloaded:
+  * Countbench from Huggingface
+  * Questions found in `data/countbench_paired_questions.json`
+  """
+
+  def _info(self) -> tfds.core.DatasetInfo:
+    """Returns the dataset metadata."""
+    features = tfds.features.FeaturesDict({
+        'image': tfds.features.Image(shape=(None, None, 3)),
+        'image_id': tfds.features.Scalar(dtype=np.int32),
+        'question': tfds.features.Text(),
+        'text': tfds.features.Text(),
+        'image_url': tfds.features.Text(),
+        'number': tfds.features.Scalar(dtype=np.int32),
+    })
+
+    return tfds.core.DatasetInfo(
+        builder=self,
+        features=features,
+        description=_DESCRIPTION,
+        supervised_keys=None,
+        homepage=_HOMEPAGE,
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Call the function which defines the splits."""
+    del dl_manager
+    return {
+        'huggingface': self._generate_examples(split='huggingface'),
+    }
+
+  def _generate_examples_hf(self):
+    """Generate examples from Huggingface parquet file.
+
+    Note that the parquet file provided on Huggingface is missing about 10%
+    of the images as can be verified by running
+    ```
+        import pyarrow.parquet as pq
+        with open(_COUNTBENCH_PARQUET_PATH, 'rb') as f:
+          x = pq.read_table(f)
+        sum([x['image'][i].is_valid for i in range(len(x['image']))])  # result: 491
+    ```
+
+    Yields:
+      An index and a dictionary with features.
+    """
+    with open(_COUNTBENCH_PARQUET_PATH, 'rb') as f:
+      df = pd.read_parquet(f)
+
+    with open(_QA_JSON_PATH, 'r') as fq:
+      df_question = pd.read_json(fq)
+
+    df['question'] = df_question
+
+    for idx, row in df.iterrows():
+      # Some entries have no image.
+      if row['image'] is None:
+        continue
+      image = np.array(PIL.Image.open(io.BytesIO(row['image']['bytes'])))
+      if len(image.shape) != 3:
+        continue  # Filter out one bad image.
+      countbenchqa_dict = {
+          'image': image,
+          'image_id': idx,
+          'question': row['question'],
+          'text': row['text'],
+          'image_url': row['image_url'],
+          'number': row['number'],
+      }
+      yield idx, countbenchqa_dict
+
+  def _generate_examples(self, split: str):
+    if split == 'huggingface':
+      yield from self._generate_examples_hf()
+    else:
+      raise ValueError(f'Unknown split: {split}')

Tipsomaly/model/big_vision/datasets/docvqa/docvqa.py

@@ -0,0 +1,110 @@
+# 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"""Implements DocVQA in TFDS structure.
+
+It's small data, so simple to run locally. First, copy the data to local disk.
+An account will be needed in https://rrc.cvc.uab.es/?ch=17&com=downloads and
+from there the task annotations and images can be fetched separatedly.
+
+  mkdir -p /tmp/data/docvqa
+  <COPY AND DECOMPRESS DOWNLOADED FILES HERE>
+
+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=docvqa
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset = tfds.load('docvqa', split='val', data_dir='/tmp/tfds')
+"""
+import json
+import os
+
+import numpy as np
+import tensorflow_datasets as tfds
+
+
+_DESCRIPTION = """DocVQA dataset."""
+
+# pylint: disable=line-too-long
+_CITATION = """
+@article{DBLP:journals/corr/abs-2007-00398,
+  author       = {Minesh Mathew and
+                  Dimosthenis Karatzas and
+                  R. Manmatha and
+                  C. V. Jawahar},
+  title        = {DocVQA: {A} Dataset for {VQA} on Document Images},
+  journal      = {CoRR},
+  volume       = {abs/2007.00398},
+  year         = {2020},
+  url          = {https://arxiv.org/abs/2007.00398},
+  eprinttype    = {arXiv},
+  eprint       = {2007.00398},
+  timestamp    = {Mon, 06 Jul 2020 15:26:01 +0200},
+  biburl       = {https://dblp.org/rec/journals/corr/abs-2007-00398.bib},
+  bibsource    = {dblp computer science bibliography, https://dblp.org}
+}
+"""
+# pylint: enable=line-too-long
+
+# When running locally (recommended), copy files as above an use these:
+_DOCVQA_PATH = '/tmp/data/docvqa/'
+
+
+class DocVQA(tfds.core.GeneratorBasedBuilder):
+  """DatasetBuilder for DocVQA 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({
+            'question_id': tfds.features.Scalar(np.int32),
+            'image/filename': tfds.features.Text(),
+            'image': tfds.features.Image(encoding_format='png'),
+            'question': tfds.features.Text(),
+            'answers': tfds.features.Sequence(tfds.features.Text()),
+        }),
+        supervised_keys=None,
+        homepage='https://www.docvqa.org/',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    return {split: self._generate_examples(split)
+            for split in ('val', 'train', 'test')}
+
+  def _generate_examples(self, split: str):
+    """Yields (key, example) tuples from split."""
+    suffix = '' if split == 'test' else '_withQT'
+    with open(os.path.join(_DOCVQA_PATH, f'{split}_v1.0{suffix}.json')) as f:
+      data = json.load(f)
+    for v in data['data']:
+      question_id = v['questionId']
+      yield question_id, {
+          'question_id': question_id,
+          'image/filename': v['image'],
+          'image': os.path.join(_DOCVQA_PATH, split, v['image']),
+          'question': v['question'],
+          'answers': v.get('answers', []),
+      }

Tipsomaly/model/big_vision/datasets/gqa/gqa.py

@@ -0,0 +1,167 @@
+# 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 GQA in a TFDS-ready structure, using Beam.
+
+Instructions below are to generate the dataset with a *local* Beam pipeline.
+It's advisable to run the Beam job on Google Cloud Dataflow, see
+  https://www.tensorflow.org/datasets/beam_datasets.
+for more details, which would significantly speed up generation. This would
+involve uploading the locally downloaded data to a GCS bucket, and then
+adding in the Beam pipeline options and your GCP/GCS bucket details
+to the `tfds build` command below (as detailed in the link).
+
+First, copy the data to local disk:
+
+  mkdir -p /tmp/data/gqa
+  wget -O /tmp/data/gqa/question1.2.zip https://downloads.cs.stanford.edu/nlp/data/gqa/questions1.2.zip?download=true
+  unzip /tmp/data/gqa/question1.2.zip
+  mv /tmp/data/gqa/question1.2/* /tmp/data/gqa/
+  wget -O /tmp/data/gqa/images.zip https://downloads.cs.stanford.edu/nlp/data/gqa/images.zip?download=true
+  unzip /tmp/data/gqa/images.zip
+
+Then, run conversion (make sure to install tensorflow-datasets for the `tfds` util):
+
+  cd big_vision/datasets
+  env TFDS_DATA_DIR=/tmp/tfds tfds build --datasets=gqa
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset = tfds.load('gqa', split='testdev_balanced', data_dir='/tmp/tfds')
+
+Some statistics:
+  train_all: 14305356 examples
+  train_balanced: 943000 examples
+  val_all: 2011853 examples
+  val_balanced: 132062 examples
+  testdev_all: 172174 examples
+  testdev_balanced: 12578 examples
+"""
+import glob
+import json
+import os
+
+import numpy as np
+import tensorflow_datasets as tfds
+
+
+_DESCRIPTION = """GQA: Visual Reasoning in the Real World."""
+
+# pylint: disable=line-too-long
+_CITATION = """
+@article{DBLP:journals/corr/abs-2306-14610,
+  author       = {Drew Hudson and
+                  Christopher Manning},
+  title        = {GQA: A New Dataset for Real-World Visual Reasoning and Compositional Question Answering},
+  journal      = {CVPR},
+  volume       = {abs/1902.09506},
+  year         = {2019},
+  url          = {https://doi.org/10.48550/arXiv.1902.09506},
+  doi          = {10.48550/arXiv.1902.09506},
+  eprinttype    = {arXiv},
+  eprint       = {1902.09506},
+  timestamp    = {Tue, 25 Jun 2019 00:00:00 +0100},
+  biburl       = {https://dblp.org/rec/journals/corr/abs-1902-09506},
+  bibsource    = {dblp computer science bibliography, https://dblp.org}
+}
+"""
+# pylint: enable=line-too-long
+
+
+_DATA_PATH = '/tmp/data/gqa/'
+
+
+class GQA(tfds.core.GeneratorBasedBuilder):
+  """DatasetBuilder for GQA 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({
+            'example_id': tfds.features.Scalar(np.int64),
+            'image/id': tfds.features.Text(),
+            'image': tfds.features.Image(encoding_format='jpeg'),
+            'question': tfds.features.Text(),
+            'answer': tfds.features.Text(),
+            'full_answer': tfds.features.Text(),
+            'is_balanced': tfds.features.Scalar(np.bool_),
+        }),
+        homepage='https://cs.stanford.edu/people/dorarad/gqa/',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    splits = [
+        # 'debug',
+        'train_all',
+        'train_balanced',
+        'testdev_all',
+        'testdev_balanced',
+        'val_all',
+        'val_balanced',
+        'challenge_all',
+        'challenge_balanced',
+    ]
+    return {split: self._generate_examples(split) for split in splits}
+
+  def _generate_examples(self, split: str):
+    """Yields (key, example) tuples from dataset."""
+    if split == 'train_all':
+      train_json_dir = os.path.join(_DATA_PATH, 'train_all_questions', '*.json')
+      json_files = glob.glob(train_json_dir)
+    else:
+      json_files = [os.path.join(_DATA_PATH, f'{split}_questions.json')]
+
+    def _prepare_data(json_path):
+      with open(os.path.join(json_path)) as f:
+        annotations = json.load(f)
+      return [(k, v) for k, v in annotations.items()]
+
+    def _process_example(entry):
+      question_id, question_data = entry
+      image_id = question_data['imageId']
+      image_path = os.path.join(_DATA_PATH, 'images', f'{image_id}.jpg')
+      answer = question_data['answer'] if 'answer' in question_data else ''
+      if 'fullAnswer' in question_data:
+        full_answer = question_data['fullAnswer']
+      else:
+        full_answer = ''
+
+      example = {
+          'example_id': question_id,
+          'image/id': image_id,
+          'image': image_path,
+          'question': question_data['question'],
+          'answer': answer,
+          'full_answer': full_answer,
+          'is_balanced': question_data['isBalanced'],
+      }
+      return question_id, example
+
+    beam = tfds.core.lazy_imports.apache_beam
+    return (
+        beam.Create(json_files)
+        | beam.FlatMap(_prepare_data)
+        | beam.Reshuffle()
+        | beam.Map(_process_example)
+    )

Tipsomaly/model/big_vision/datasets/infovqa/infovqa.py

@@ -0,0 +1,141 @@
+# 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"""Implements InfoVqa in TFDS structure.
+
+First, download and unzip the dataset from https://rrc.cvc.uab.es/?ch=17
+and place it in /tmp/data/infovqa.
+
+Then, run conversion locally (make sure to install tensorflow-datasets for the `tfds` util):
+
+  cd third_party/py/big_vision/datasets
+  env TFDS_DATA_DIR=/tmp/tfds tfds build --datasets=infovqa
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset = tfds.load('infovqa', split='train', data_dir='/tmp/tfds')
+
+Dataset splits:
+  train: 23946 examples/questions (4406 images)
+  val: 2801 examples/questions (500 images)
+  test: 3288 examples/questions (579 images) (no answers)
+
+Recommended training splits:
+  train: train[:95%] (22749 examples/questions)
+  minitrain: train[:5%] (1197 examples/questions)
+  minival: train[95%:] (1197 examples/questions)
+  eval: val (2801 examples/questions)
+
+Note that according to task description in
+https://rrc.cvc.uab.es/?ch=17&com=tasks:
+  - Order of items in a multi span answer does not matter. Therefore, we include
+    all permutations of the answer in the val split.
+  - Answers are not case sensitive. We leave it to the user to lower case
+    answers if they want to.
+"""
+import itertools
+import json
+import os
+
+import numpy as np
+import tensorflow_datasets as tfds
+
+
+_DESCRIPTION = """InfographicVQA dataset."""
+
+# pylint: disable=line-too-long
+_CITATION = """
+@inproceedings{Mathew_2022,
+   title={InfographicVQA},
+   url={http://dx.doi.org/10.1109/WACV51458.2022.00264},
+   DOI={10.1109/wacv51458.2022.00264},
+   booktitle={2022 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV)},
+   publisher={IEEE},
+   author={Mathew, Minesh and Bagal, Viraj and Tito, Ruben and Karatzas, Dimosthenis and Valveny, Ernest and Jawahar, C. V.},
+   year={2022},
+   month=jan }
+"""
+# pylint: enable=line-too-long
+
+# When running locally (recommended), copy files as above an use these:
+_INFOVQA_PATH = '/tmp/data/infovqa/'
+_ANNOTATIONS = {
+    'train': 'infographicsVQA_train_v1.0.json',
+    'val': 'infographicsVQA_val_v1.0_withQT.json',
+    'test': 'infographicsVQA_test_v1.0.json',
+    }
+
+
+class Infovqa(tfds.core.GeneratorBasedBuilder):
+  """DatasetBuilder for infovqa dataset."""
+
+  VERSION = tfds.core.Version('1.1.0')
+  RELEASE_NOTES = {
+      '1.0.0': 'First release.',
+      '1.1.0': 'Add multi-span permutations to the val split answers.',
+      }
+
+  def _info(self):
+    """Returns the metadata."""
+
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=_DESCRIPTION,
+        features=tfds.features.FeaturesDict({
+            'question_id': tfds.features.Scalar(np.int32),
+            'filename': tfds.features.Text(),
+            'image': tfds.features.Image(encoding_format='jpeg'),
+            'question': tfds.features.Text(),
+            'answers': tfds.features.Sequence(tfds.features.Text()),
+        }),
+        supervised_keys=None,
+        homepage='https://www.docvqa.org/datasets/infographicvqa',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    return {split: self._generate_examples(split)
+            for split in ('train', 'val', 'test')}
+
+  def _generate_examples(self, split):
+    """Yields (key, example) tuples from test set."""
+    annot_fname = os.path.join(_INFOVQA_PATH, _ANNOTATIONS[split])
+    with open(annot_fname, 'r') as f:
+      data = json.loads(f.read())
+
+    for x in data['data']:
+      yield x['questionId'], {
+          'question_id': x['questionId'],
+          'filename': x['image_local_name'],
+          'image': os.path.join(_INFOVQA_PATH, 'images', x['image_local_name']),
+          'question': x['question'],
+          'answers': maybe_permute(x.get('answers', []), split),
+      }
+
+
+def maybe_permute(answers, split):
+  if split != 'val':
+    return answers
+  new_answers = []
+  for x in answers:
+    if ', ' in x:  # Create all permutations.
+      # The first element remains the same.
+      new_answers.extend([', '.join(y)
+                          for y in itertools.permutations(x.split(', '))])
+    else:
+      new_answers.append(x)
+  return new_answers

Tipsomaly/model/big_vision/datasets/okvqa/okvqa.py

@@ -0,0 +1,213 @@
+# 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"""Implements the OKVQA dataset for TFDS.
+
+Download the required files from https://okvqa.allenai.org/download.html:
+
+mkdir -p /tmp/tfds
+cd /tmp/tfds/
+wget http://images.cocodataset.org/zips/train2014.zip
+wget http://images.cocodataset.org/zips/val2014.zip
+wget https://okvqa.allenai.org/static/data/mscoco_train2014_annotations.json.zip
+wget https://okvqa.allenai.org/static/data/mscoco_val2014_annotations.json.zip
+wget https://okvqa.allenai.org/static/data/OpenEnded_mscoco_train2014_questions.json.zip
+wget https://okvqa.allenai.org/static/data/OpenEnded_mscoco_val2014_questions.json.zip
+unzip val2014.zip
+unzip train2014.zip
+unzip OpenEnded_mscoco_train2014_questions.json.zip
+unzip OpenEnded_mscoco_val2014_questions.json.zip
+unzip mscoco_train2014_annotations.json.zip
+unzip mscoco_val2014_annotations.json.zip
+
+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=okvqa
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset = tfds.load('okvqa', split='val', data_dir='/tmp/tfds')
+"""
+
+import json
+import os
+from typing import Any
+import numpy as np
+import tensorflow_datasets as tfds
+
+_DESCRIPTION = """
+OKVQA addresses the task of VQA with outside knowledge.
+This version of the dataset contains:
+- Questions + Answers from OKVQA.
+- Images from COCO.
+"""
+
+_CITATION = """
+@InProceedings{okvqa,
+author = {Kenneth Marino and Mohammad Rastegari and Ali Farhadi and Roozbeh Mottaghi},
+title = {OK-VQA: A Visual Question Answering Benchmark Requiring External Knowledge},
+booktitle = {Conference on Computer Vision and Pattern Recognition (CVPR)},
+year = {2019},
+}
+"""
+
+ANNOTATION_FILE = {
+    'train': 'mscoco_train2014_annotations.json',
+    'val': 'mscoco_val2014_annotations.json',
+}
+QUESTIONS_FILE = {
+    'train': 'OpenEnded_mscoco_train2014_questions.json',
+    'val': 'OpenEnded_mscoco_val2014_questions.json',
+}
+QUESTION_TYPES = {
+    'one': 'Vehicles and Transportation',
+    'two': 'Brands, Companies and Products',
+    'three': 'Objects, Material and Clothing',
+    'four': 'Sports and Recreation',
+    'five': 'Cooking and Food',
+    'six': 'Geography, History, Language and Culture',
+    'seven': 'People and Everyday life',
+    'eight': 'Plants and Animals',
+    'nine': 'Science and Technology',
+    'ten': 'Weather and Climate',
+    'other': 'Other',
+}
+
+
+# When running locally (recommended), copy files as above an use these:
+_OKVQA_PATH = '/media/scratch/okvqa'
+
+
+class OkVqa(tfds.core.GeneratorBasedBuilder):
+  """Import COCO dataset for OKVQA with KAT features."""
+
+  VERSION = tfds.core.Version('1.0.0')
+  RELEASE_NOTES = {'1.0.0': 'Changed to array record format.'}
+  MANUAL_DOWNLOAD_INSTRUCTIONS = """
+  In manual_dir/ you should have a directory okvqa which contains the
+  following files and directories:
+  From the OKVQA dataset:
+  - mscoco_train2014_annotations.json
+  - mscoco_val2014_annotations.json
+  - OpenEnded_mscoco_train2014_questions.json
+  - OpenEnded_mscoco_val2014_questions.json
+  - train2014.zip
+  - val2014.zip
+  """
+
+  def _info(self) -> tfds.core.DatasetInfo:
+    """Returns the dataset metadata."""
+    features = tfds.features.FeaturesDict({
+        'image': tfds.features.Image(shape=(None, None, 3)),
+        'image_id': tfds.features.Scalar(dtype=np.int64),
+        'answer_type': tfds.features.Text(),
+        'answers': tfds.features.Sequence(tfds.features.Text()),
+        'answers_confidence': tfds.features.Tensor(shape=[10], dtype=np.bool_),
+        'answers_raw': tfds.features.Sequence(tfds.features.Text()),
+        'question_id': tfds.features.Scalar(dtype=np.int64),
+        'question_type': tfds.features.Text(),
+        'question_type_readable': tfds.features.Text(),
+        'question': tfds.features.Text(),
+    })
+
+    return tfds.core.DatasetInfo(
+        builder=self,
+        features=features,
+        description=_DESCRIPTION,
+        supervised_keys=None,
+        homepage='https://okvqa.allenai.org/',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager) -> ...:
+    """Call the function which defines the splits."""
+    # data_dir = dl_manager.manual_dir
+    data_dir = _OKVQA_PATH
+    return {
+        'train': self._generate_examples(data_dir, 'train'),
+        'val': self._generate_examples(data_dir, 'val'),
+    }
+
+  def _generate_examples(self, data_dir: str, split: str) -> ...:
+    annotations = get_okvqa_annotations(data_dir, split)
+
+    for question_id, annotation in annotations.items():
+      image_id = annotation['image_id']
+
+      # Sanity check.
+      if len(annotation['answers']) != 10:
+        num_answers = len(annotation['answers'])
+        raise ValueError(
+            f'The number of answers for {image_id} is not 10 but {num_answers}')
+
+      feature_dict = {
+          'image': self.get_image_path(data_dir, split, image_id),
+          'image_id': image_id,
+          'answer_type': annotation['answer_type'],
+          'answers': [a['answer'] for a in annotation['answers']],
+          'answers_confidence': _get_answer_confidence(annotation['answers']),
+          'answers_raw': [a['raw_answer'] for a in annotation['answers']],
+          'question_id': annotation['question_id'],
+          'question_type': annotation['question_type'],
+          'question_type_readable': QUESTION_TYPES[annotation['question_type']],
+          'question': annotation['question'],
+      }
+      yield f'{question_id}', feature_dict
+
+  def get_image_path(self, data_dir: str, split: str, image_id: int) -> str:
+    subdir = {'train': 'train2014', 'val': 'val2014'}[split]
+    return f'{data_dir}/{subdir}/COCO_{subdir}_{image_id:012d}.jpg'
+
+
+def _get_answer_confidence(answers: list[dict[str, str]]) -> np.ndarray:
+  """Get OKVQA answer confidences as bool."""
+  confidences = []
+  for a in answers:
+    confidence = a['answer_confidence']
+    if confidence == 'yes':
+      confidences.append(True)
+    elif confidence == 'no':
+      confidences.append(False)
+    else:
+      raise ValueError(f'Unknown confidence: {confidence}')
+  return np.array(confidences, dtype=bool)
+
+
+def _read_json(
+    data_dir: str, file: str, key: str
+) -> dict[int, dict[str, Any]]:
+  with open(os.path.join(data_dir, file)) as f:
+    data = json.load(f)
+  questions = {d['question_id']: d for d in data[key]}
+  return questions
+
+
+def get_okvqa_annotations(
+    data_dir: str, split: str
+) -> dict[int, dict[str, Any]]:
+  """Return okvqa annotations (quesions and answers) as dictionary."""
+  questions = _read_json(data_dir, QUESTIONS_FILE[split], 'questions')
+  annotations = _read_json(data_dir, ANNOTATION_FILE[split], 'annotations')
+
+  assert len(annotations) == len(questions)
+  for question_id, question in questions.items():
+    assert question['image_id'] == annotations[question_id]['image_id']
+    assert question['question_id'] == annotations[question_id]['question_id']
+    annotations[question_id]['question'] = question['question']
+
+  return annotations

Tipsomaly/model/big_vision/datasets/pope/pope.py

@@ -0,0 +1,145 @@
+# 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"""Implements POPE test-set in TFDS structure.
+
+It's small data, so simple to run locally. First, copy the data to local disk:
+First download json files from https://github.com/AoiDragon/POPE; then download
+MSCOCO (val 2014) images from https://cocodataset.org/#download
+
+  mkdir -p /tmp/data/pope/
+  mkdir -p /tmp/data/pope/pope/
+  mkdir -p /tmp/data/pope/images/
+  git clone https://github.com/AoiDragon/POPE.git
+  cp POPE/output/coco/* /tmp/data/pope/pope/
+  wget http://images.cocodataset.org/zips/val2014.zip
+  unzip val2014.zip
+  cp -r val2014/ /tmp/data/pope/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=pope
+
+Example to load:
+
+  import tensorflow_datasets as tfds
+  dataset_random = tfds.load('pope/pope_random', split='test', data_dir='/tmp/tfds')
+  dataset_popular = tfds.load('pope/pope_popular', split='test', data_dir='/tmp/tfds')
+  dataset_adversarial = tfds.load('pope/pope_adversarial', split='test', data_dir='/tmp/tfds')
+
+"""
+import json
+import os
+
+import numpy as np
+import tensorflow_datasets as tfds
+
+
+_DESCRIPTION = """POPE dataset."""
+
+# pylint: disable=line-too-long
+_CITATION = """
+@inproceedings{li-etal-2023-evaluating,
+    title = "Evaluating Object Hallucination in Large Vision-Language Models",
+    author = "Li, Yifan  and
+      Du, Yifan  and
+      Zhou, Kun  and
+      Wang, Jinpeng  and
+      Zhao, Xin  and
+      Wen, Ji-Rong",
+    editor = "Bouamor, Houda  and
+      Pino, Juan  and
+      Bali, Kalika",
+    booktitle = "Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing",
+    month = dec,
+    year = "2023",
+    address = "Singapore",
+    publisher = "Association for Computational Linguistics",
+    url = "https://aclanthology.org/2023.emnlp-main.20",
+    doi = "10.18653/v1/2023.emnlp-main.20",
+    pages = "292--305",
+    abstract = "Inspired by the superior language abilities of large language models (LLM), large vision-language models (LVLM) have been recently proposed by integrating powerful LLMs for improving the performance on complex multimodal tasks. Despite the promising progress on LVLMs, we find that they suffer from object hallucinations, i.e., they tend to generate objects inconsistent with the target images in the descriptions. To investigate it, this work presents the first systematic study on object hallucination of LVLMs. We conduct the evaluation experiments on several representative LVLMs, and show that they mostly suffer from severe object hallucination issues. We further discuss that the visual instructions may influence the hallucination, and find that: objects that frequently appear in the visual instructions or co-occur with the image objects are obviously prone to be hallucinated by LVLMs. Besides, we further design a polling-based query method called POPE for better evaluation of object hallucination. Experiment results show that our POPE can evaluate object hallucination in a more stable and flexible way.",
+}
+"""
+# pylint: enable=line-too-long
+
+# When running locally (recommended), copy files as above and use these:
+_POPE_PATH = '/tmp/data/pope/'
+
+
+class POPEConfig(tfds.core.BuilderConfig):
+  """Configuration to build the dataset."""
+
+  pass
+
+
+class POPE(tfds.core.GeneratorBasedBuilder):
+  """DatasetBuilder for POPE dataset."""
+
+  VERSION = tfds.core.Version('1.0.0')
+  RELEASE_NOTES = {'1.0.0': 'First release.'}
+  BUILDER_CONFIGS = [
+      POPEConfig(name='pope_random', description='Random set'),
+      POPEConfig(name='pope_popular', description='Popular set'),
+      POPEConfig(name='pope_adversarial', description='Adversarial set'),
+  ]
+
+  def _info(self):
+    """Returns the metadata."""
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=_DESCRIPTION,
+        features=tfds.features.FeaturesDict({
+            'question_id': tfds.features.Scalar(np.int32),
+            'image/filename': tfds.features.Text(),
+            'image': tfds.features.Image(encoding_format='png'),
+            'question': tfds.features.Text(),
+            'answer': tfds.features.Text(),
+            'thing': tfds.features.Text(),
+        }),
+        supervised_keys=None,
+        homepage='https://github.com/AoiDragon/POPE',
+        citation=_CITATION,
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    return {'test': self._generate_examples('test', self.builder_config.name)}
+
+  def _generate_examples(self, split: str, source: str):
+    """Yields (key, example) tuples from test set."""
+    annot_fname = os.path.join(
+        _POPE_PATH, f'pope/coco_{source}.json'
+    )
+
+    with open(annot_fname, 'r') as f:
+      data = [json.loads(line) for line in f]
+
+    for idx, v in enumerate(data):
+      question = v['text']
+      thing = (
+          question.replace('Is there an ', '')
+          .replace('Is there a ', '')
+          .replace(' in the image?', '')
+      )
+      yield idx, {
+          'question_id': idx,
+          'image/filename': v['image'],
+          'image': os.path.join(_POPE_PATH, 'images/val2014/', v['image']),
+          'question': question,
+          'answer': v['label'],
+          'thing': thing,
+      }

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/coco_caption.py

@@ -0,0 +1,145 @@
+# 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.
+
+"""Evaluator for caption generation metrics used for the MS COCO dataset."""
+import collections
+import functools
+import os
+import tempfile
+
+import big_vision.evaluators.common as c
+import big_vision.input_pipeline
+import big_vision.pp.builder
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+
+from pycocoevalcap.bleu import bleu
+from pycocoevalcap.cider import cider
+from pycocoevalcap.meteor import meteor
+from pycocoevalcap.rouge import rouge
+from pycocoevalcap.spice import spice
+from pycocoevalcap.tokenizer import ptbtokenizer
+
+import jax
+
+from tensorflow.io import gfile
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+
+
+class Evaluator:
+  """Evaluator for caption generation metrics used for the MS COCO dataset.
+
+  See https://arxiv.org/pdf/1504.00325.pdf or the repository implementing it
+  https://github.com/tylin/coco-caption for details on the metrics. This code
+  uses the python3 pip package from: https://github.com/salaniz/pycocoevalcap
+
+  Note that both the model caption and the ground truth reference captions are
+  further processed with the PTBTokenizer before computing scores.
+
+  `predict_fn` accepts arbitrary dictionaries of parameters and data, where
+  the data dictionary is produced by the `pp_fn` op. It is expected to output a
+  dict containing tokenized captions.
+
+  `pp_fn` must have fields: "image/id" and "captions".
+  """
+
+  def __init__(
+      self, predict_fn, tokenizer=None,
+      metrics=("cider",),  # Default to only cider. We often just look at that.
+      preds_outfile="{workdir}/{name}_{split}_preds.json",
+      annot_outfile="{workdir}/{name}_{split}_annotations.json",
+      *, data, devices, **kw
+  ):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={"image/id", "captions"}, data=data, devices=devices, **kw)
+
+    self.preds_outfile = c.resolve_outfile(
+        preds_outfile, name=data.get("name"), split=data.get("split"))
+    self.annot_outfile = c.resolve_outfile(
+        annot_outfile, name=data.get("name"), split=data.get("split"))
+
+    self.metrics = metrics
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token)
+
+  def run(self, train_state):
+    """Run eval."""
+    gts = []
+    res = []
+
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded generated tokens.
+      tokens = self.decode(train_state, batch)
+
+      # (local_batch,)
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      image_ids = batch["image/id"][ex_masks]
+      pred_captions = self.tok.to_str(tokens[ex_masks])
+
+      for image_id, caption in zip(image_ids, pred_captions):
+        res.append({"image_id": image_id.item(), "caption": caption})
+
+      for image_id, captions in zip(image_ids, batch["captions"]):
+        for caption in captions:
+          gts.append({"image_id": image_id.item(), "caption": caption.item()})
+
+    # Write model outputs following: https://cocodataset.org/#format-results
+    # Use same format for gt although that is not the usual format for them.
+    res = c.multiprocess_write_json(self.preds_outfile, res)
+    gts = c.multiprocess_write_json(self.annot_outfile, gts)
+
+    if jax.process_index():  # Host0 gets all preds and does eval.
+      return
+
+    outs = self.evaluate(gts, res)
+    for key, score in outs.items():
+      yield key, score
+
+  def evaluate(self, gt_annotations, res_annotations):
+    """Creates scorers and run evaluation."""
+    scorers = {
+        "rouge": rouge.Rouge,
+        "cider": cider.Cider,
+        "bleu-4": bleu.Bleu,
+        "spice": spice.Spice,
+        "meteor": meteor.Meteor,
+    }
+
+    # Reformat gts and res from [{"image_id": int|str, "caption": str}] to
+    # {int_image_id: [{"caption": str}]} as expected by tokenizer and scorers.
+    # Note there are multiple reference captions for the ground truth but only
+    # one for the model predictions.
+    iid_map = collections.defaultdict(lambda: len(iid_map))
+    res = {iid_map[x["image_id"]]: [x] for x in res_annotations}
+    gts = collections.defaultdict(list)
+    for x in gt_annotations:
+      gts[iid_map[x["image_id"]]].append(x)
+    assert sorted(gts.keys()) == sorted(res.keys())
+
+    # Tokenize captions and predictions using coco tokenizer.
+    coco_tokenizer = ptbtokenizer.PTBTokenizer()
+    gts = coco_tokenizer.tokenize(gts)
+    res = coco_tokenizer.tokenize(res)
+
+    scores = {}
+    for metric in self.metrics:
+      scorer = scorers[metric]()
+      scores[metric], _ = scorer.compute_score(gts, res)
+    return scores

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/rsvqa.py

@@ -0,0 +1,173 @@
+# 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.
+
+"""Evaluator for simple VQA variants with per answer-type metrics.
+
+According to the (A-)OKVAQ papers, the eval for these datasets should follow
+VQAv2. But here we don't track different answer-types, and don't do any
+leave-one-out averaging, as this isn't done in the official implementation at
+https://github.com/allenai/aokvqa/blob/main/evaluation/eval_predictions.py
+either.
+"""
+
+import functools
+
+import big_vision.evaluators.common as c
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+import editdistance
+
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+
+QUESTION_TYPES = ("comp", "count", "presence", "rural_urban", "area")
+
+ACC_SUBSETS = (
+    ("nonum", ("comp", "presence", "rural_urban")),  # rsvqa_lr
+    ("nonum", ("comp", "presence")),  # rsvqa_hr
+)
+
+
+class Evaluator:
+  """Evaluator for simple VQA tasks."""
+
+  def __init__(
+      self, predict_fn, tokenizer, to_lower=False,
+      outfile="{workdir}/{split}.json",
+      *, data, devices, **kw):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={"answers", "answer", "question_id", "question_type"},
+        data=data, devices=devices, **kw)
+
+    self.outfile = c.resolve_outfile(outfile, split=data.get("split"))
+
+    # We'll need the tokenizer to detokenize the model outputs later.
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.postproc = (lambda s: s.lower()) if to_lower else lambda s: s
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token)
+
+  def run(self, train_state):
+    """Does one evaluation run, yields metrics."""
+
+    accuracies = []
+    accuracies_any = []
+    counts_per_type = {t: 0 for t in QUESTION_TYPES}
+    accuracies_per_type = {t: [] for t in QUESTION_TYPES}
+    anls_values = []
+    json_out = []
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded generated tokens.
+      tokens = self.decode(train_state, batch)  # (B,L,E)
+
+      # (local_batch,) that indicates padding examples (0) vs real examples (1).
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      # Turn predictions into texts and then scores, one by one.
+      for i in range(len(tokens)):
+        if ex_masks[i] == 0:  # Skip last-batch padding examples
+          continue
+
+        answer = self.postproc(self.tok.to_str(tokens[i], stop_at_eos=True))
+
+        # Now we have two commonly used VQA evaluation modes:
+        if "answer" in batch:
+          # single GT (eg ocrvqa): just compare to that answer, done.
+          gt = self.postproc(batch["answer"][i])
+          gts = [gt]
+          accuracies.append(float(answer == gt))
+          accuracies_any.append(float(answer == gt))
+          anls_values.append(anls_metric(gt, answer))
+        elif "answers" in batch and (gt_answers := batch["answers"][i]).size:
+          # multiple GTs (eg okvqa): introduced by VQA, compare to each of them
+          # with a threshold, see also: https://visualqa.org/evaluation.html
+          gts = [self.postproc(a) for a in gt_answers]
+          num_match = sum([answer == gt for gt in gts])
+          accuracies.append(min(1.0, num_match / 3.0))
+          accuracies_any.append(min(1.0, float(num_match)))
+          anls_values.append(max(anls_metric(gt, answer) for gt in gts))
+          accuracies_per_type[batch["question_type"][i]].append(
+              accuracies_any[-1]
+          )
+          counts_per_type[batch["question_type"][i]] += 1
+        else:
+          gts = []
+
+        json_out.append({
+            "question_id": batch["question_id"][i].item(),
+            "answer": answer} | ({"gts": gts} if gts else {}))
+
+    # At this point `accuracies` is a list of per-example scores. However,
+    # remember that each host holds a different subset of the examples! So if
+    # we were to just return the mean accuracy here, we would effectively only
+    # have evaluated on the main host's (who writes metrics) subset!
+    # So now, we need to compute global means.
+    # There is one more caveat: `process_sum` needs the summands on each host
+    # to have the same size. So we either need to include dummy values for
+    # the padding examples (last batch, annoying), or we only sum scalars as in
+    # sufficient statistics, which we do here.
+    sum_accs, sum_accs_any, sum_anls, num_accs, num = c.process_sum(
+        [sum(accuracies), sum(accuracies_any), sum(anls_values),
+         len(accuracies), len(json_out)])
+
+    sum_accs_per_type, sum_cnts_per_type = c.process_sum(
+        [{k: sum(v) for k, v in accuracies_per_type.items()}, counts_per_type]
+    )
+
+    # Yielding metric_name, value means logging the metric.
+    if num_accs:
+      yield "acc", sum_accs / num_accs
+      yield "acc_any", sum_accs_any / num_accs  # Overall Accuracy (OA).
+      yield "anls", sum_anls / num_accs
+      acc_types = {}
+      for k, v in sum_accs_per_type.items():
+        if sum_cnts_per_type[k]:
+          acc_types[k] = v / sum_cnts_per_type[k]
+          yield f"acc_{k}", acc_types[k]
+      yield "acc_avg", sum(acc_types.values()) / len(acc_types)  # Avg acc (AA).
+      for postfix, types in ACC_SUBSETS:
+        if all(t in acc_types for t in types):
+          yield f"acc_avg_{postfix}", sum(
+              [v for k, v in acc_types.items() if k in types]
+          ) / len(types)  # Average accuracy per question types subset.
+    yield "num", num  # Just for sanity checks.
+    c.multiprocess_write_json(self.outfile, json_out)
+
+
+def anls_metric(target: str, prediction: str, theta: float = 0.5):
+  """Calculates ANLS for DocVQA.
+
+  There does not seem to be an official evaluation script.
+  Public implementation on which this implementation is based:
+  https://github.com/herobd/layoutlmv2/blob/main/eval_docvqa.py#L92
+
+  Original paper (see Eq 1): https://arxiv.org/pdf/1907.00490.pdf
+
+  Args:
+    target: Target string.
+    prediction: Predicted string.
+    theta: Filter threshold set to 0.5 for DocVQA.
+
+  Returns:
+    ANLS score.
+  """
+  if target:
+    edit_distance = editdistance.eval(target, prediction)
+    normalized_ld = edit_distance / max(len(target), len(prediction))
+    return 1 - normalized_ld if normalized_ld < theta else 0
+  else:
+    return float(prediction == "")

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/science_qa.py

@@ -0,0 +1,122 @@
+# 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.
+
+"""Evaluator for ScienceQA.
+
+based on the official implementation at
+https://github.com/lupantech/ScienceQA/blob/main/models/run_gpt3.py
+"""
+
+import functools
+import re
+
+import big_vision.evaluators.common as c
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+FAILURE = "failed"
+
+
+class Evaluator:
+  """Evaluator for simple VQA tasks."""
+
+  def __init__(
+      self, predict_fn, tokenizer,
+      outfile="{workdir}/{split}.json",
+      out_question_key="question_id",
+      *, data, devices, **kw):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={"answer", "question_id"}, data=data, devices=devices, **kw)
+
+    self.outfile = c.resolve_outfile(outfile, split=data.get("split"))
+    self.out_question_key = out_question_key
+
+    # We'll need the tokenizer to detokenize the model outputs later.
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token
+    )
+
+  def postproc(self, raw_answer):
+    """Post-processes the raw answer. extract a, b, c from the string."""
+    match = re.match(
+        pattern=r"the answer is ([a-z])\.", string=raw_answer.lower()
+    )
+    if match:
+      return match.groups()[0]  # 'a', 'b', ...
+    else:
+      return FAILURE
+
+  def run(self, train_state):
+    """Does one evaluation run, yields metrics."""
+
+    accuracies = []
+    fail_parse = []
+    json_out = []
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded generated tokens.
+      tokens = self.decode(train_state, batch)
+
+      # (local_batch,) that indicates padding examples (0) vs real examples (1).
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      # Turn predictions into texts and then scores, one by one.
+      for i in range(len(tokens)):
+        if ex_masks[i] == 0:  # Skip last-batch padding examples
+          continue
+
+        raw_answer = self.tok.to_str(tokens[i], stop_at_eos=True)
+        answer = self.postproc(raw_answer)
+        if "answer" in batch:
+          gt = self.postproc(batch["answer"][i])
+          gts = [gt]
+          accuracies.append(float(answer == gt))
+          fail_parse.append(float(answer == FAILURE))
+        else:
+          gts = []
+
+        json_out.append(
+            {
+                self.out_question_key: batch["question_id"][i].item(),
+                "raw_answer": raw_answer,
+                "answer": answer,
+            }
+            | ({"gts": gts} if gts else {})
+        )
+
+    # At this point `accuracies` is a list of per-example scores. However,
+    # remember that each host holds a different subset of the examples! So if
+    # we were to just return the mean accuracy here, we would effectively only
+    # have evaluated on the main host's (who writes metrics) subset!
+    # So now, we need to compute global means.
+    # There is one more caveat: `process_sum` needs the summands on each host
+    # to have the same size. So we either need to include dummy values for
+    # the padding examples (last batch, annoying), or we only sum scalars as in
+    # sufficient statistics, which we do here.
+    sum_accs, num_parsefail, num_accs, num = c.process_sum(
+        [sum(accuracies), sum(fail_parse), len(accuracies), len(json_out)]
+    )
+
+    # Yielding metric_name, value means logging the metric.
+    if num_accs > 0:
+      yield "acc", sum_accs / num_accs
+      yield "parsefail", num_parsefail / num_accs
+
+    yield "num", num  # Just for sanity checks.
+    c.multiprocess_write_json(self.outfile, json_out)

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/segmentation.py

@@ -0,0 +1,270 @@
+# 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.
+
+"""Evaluator for segmentation."""
+
+import functools
+
+import big_vision.evaluators.common as c
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import numpy as np
+import PIL.Image
+
+from tensorflow.io import gfile
+
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = 'jit'
+
+
+def _inrange(a, min_value, max_value):
+  return (np.clip(a, min_value, max_value) == a).all()
+
+
+def _area(y1, x1, y2, x2):
+  return max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
+
+
+class Evaluator:
+  """Evaluator for instance segmentation."""
+
+  def __init__(self, predict_fn, tokenizer,
+               model='oi', det_ious=(0.5, 0.75),
+               *, devices, **kw):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={'prefix', 'suffix', 'objects/mask', 'objects/bbox'},
+        devices=devices, **kw)
+
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token)
+    tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.loc0 = np.array(tok.to_int('<loc0000>'))
+    self.seg0 = np.array(tok.to_int('<seg000>'))
+    # Verify tokenizer has `tokensets=("loc", "seg")`
+    assert self.loc0.shape == (1,), self.loc0
+    assert self.seg0.shape == (1,), self.seg0
+    self.reconstruct_masks = get_reconstruct_masks(model)
+    self.det_ious = det_ious
+
+  def run(self, train_state):
+    """Does one evaluation run, yields metrics."""
+    ious = []  # NOTE: no point to split in s/m/l: all objects are L (>96px²)
+    det_by_iou = {iou: [] for iou in self.det_ious}
+    invalid = total = 0
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+
+      decoded = self.decode(train_state, batch)
+
+      not_padding = u.get_local_slice_from_fsarray(batch['_mask'])
+      decoded = u.get_local_slice_from_fsarray(decoded)[not_padding]
+
+      # Note, gt masks are in full original image resolution.
+      gt_masks = [gt[:, :, 0] > 0 for gt in batch['objects/mask'][not_padding]]
+      gt_bbs = [gt for gt in batch['objects/bbox'][not_padding]]
+
+      valid = []
+      tokens = np.zeros([decoded.shape[0], 4 + 16], np.int32)
+      for i, dec in enumerate(decoded):
+        # TODO: b/andstein - do we need to optimize this loop?
+        t = np.r_[dec[:4] - self.loc0, dec[4:4 + 16] - self.seg0]  # Ignore rest
+        if (
+            len(t) == 4 + 16  # Full prediction
+            and _inrange(t[:4], 0, 1023)  # Valid box tokens
+            and _inrange(t[4:], 0, 127)  # Valid seg tokens
+            and t[2] > t[0] and t[3] > t[1]  # Valid box
+        ):
+          valid.append(True)
+          tokens[i] = t
+        else:
+          valid.append(False)
+
+      tocpu = lambda x: jax.device_put(x, jax.local_devices(backend='cpu')[0])
+      seg_indices = np.array(tokens[:, 4:])
+      mask64 = jax.device_get(self.reconstruct_masks(tocpu(seg_indices)))
+      mask64 = mask64[..., 0]
+      bbox = tokens[:, :4] / 1023  # Back to [0.0 ... 1.0]
+
+      for v, m64, gtm, bb, gtbb in zip(valid, mask64, gt_masks, bbox, gt_bbs):
+        # TODO: b/andstein - do we need to optimize this loop?
+        total += 1
+        h, w = gtm.shape  # gt is full/original image resolution mask.
+
+        # First, compute detection iou, in [0.0 ... 1.0] coordinate space.
+        y1, x1, y2, x2 = bb
+        gty1, gtx1, gty2, gtx2 = gtbb
+        ibb = max(y1, gty1), max(x1, gtx1), min(y2, gty2), min(x2, gtx2)
+        box_iou = _area(*ibb) / (_area(*bb) + _area(*gtbb) - _area(*ibb))
+        for iou_thresh in det_by_iou:
+          det_by_iou[iou_thresh].append(iou_thresh <= box_iou)
+
+        # Next, we convert to pixel coordinates and compute mask iou.
+        gt_area = gtm.sum()
+        y1, x1, y2, x2 = map(int, (y1 * h, x1 * w, y2 * h, x2 * w))
+
+        # Avoid compute-intensive mask stuff for invalid preds:
+        if not v or x2 <= x1 or y2 <= y1:  # Can still happen after int().
+          iou = 0.0
+          invalid += 1
+        else:
+          mi = np.asarray(
+              PIL.Image.fromarray(m64).resize(  # pytype: disable=wrong-arg-types  # pillow-102-upgrade
+                  [x2 - x1, y2 - y1], resample=PIL.Image.BILINEAR  # pytype: disable=module-attr
+              )
+          )  # Predicted mask in box-sized image.
+          mi = mi > 0.0  # Mask decoder output in [-1.0 ... 1.0]
+          iarea = (gtm[y1:y2, x1:x2] & mi).sum()  # Intersection pixels.
+          iou = iarea / (gt_area + mi.sum() - iarea)
+        ious.append(iou)
+
+    # Done going over all batches, now collect results from all processes.
+    sum_ious, num_ious, sum_dets, num_dets, num_invalid, num = c.process_sum([
+        sum(ious), len(ious),
+        {k: sum(v) for k, v in det_by_iou.items()},
+        {k: len(v) for k, v in det_by_iou.items()},
+        invalid, total
+    ])
+
+    yield 'miou', sum_ious / num_ious
+    for k in sum_dets:
+      yield f'boxacc/{k}', sum_dets[k] / num_dets[k]
+    yield 'invalid', num_invalid
+    yield 'total', num
+
+
+_KNOWN_MODELS = {
+    # Trained on open images.
+    'oi': 'gs://big_vision/paligemma/vae-oid.npz',
+}
+
+
+def _get_params(checkpoint):
+  """Converts PyTorch checkpoint to Flax params."""
+
+  def transp(kernel):
+    return np.transpose(kernel, (2, 3, 1, 0))
+
+  def conv(name):
+    return {
+        'bias': checkpoint[name + '.bias'],
+        'kernel': transp(checkpoint[name + '.weight']),
+    }
+
+  def resblock(name):
+    return {
+        'Conv_0': conv(name + '.0'),
+        'Conv_1': conv(name + '.2'),
+        'Conv_2': conv(name + '.4'),
+    }
+
+  return {
+      '_embeddings': checkpoint['_vq_vae._embedding'],
+      'Conv_0': conv('decoder.0'),
+      'ResBlock_0': resblock('decoder.2.net'),
+      'ResBlock_1': resblock('decoder.3.net'),
+      'ConvTranspose_0': conv('decoder.4'),
+      'ConvTranspose_1': conv('decoder.6'),
+      'ConvTranspose_2': conv('decoder.8'),
+      'ConvTranspose_3': conv('decoder.10'),
+      'Conv_1': conv('decoder.12'),
+  }
+
+
+def _quantized_values_from_codebook_indices(codebook_indices, embeddings):
+  batch_size, num_tokens = codebook_indices.shape
+  assert num_tokens == 16, codebook_indices.shape
+  unused_num_embeddings, embedding_dim = embeddings.shape
+
+  encodings = jnp.take(embeddings, codebook_indices.reshape((-1)), axis=0)
+  encodings = encodings.reshape((batch_size, 4, 4, embedding_dim))
+  return encodings
+
+
+class ResBlock(nn.Module):
+  features: int
+
+  @nn.compact
+  def __call__(self, x):
+    original_x = x
+    x = nn.Conv(features=self.features, kernel_size=(3, 3), padding=1)(x)
+    x = nn.relu(x)
+    x = nn.Conv(features=self.features, kernel_size=(3, 3), padding=1)(x)
+    x = nn.relu(x)
+    x = nn.Conv(features=self.features, kernel_size=(1, 1), padding=0)(x)
+    return x + original_x
+
+
+class Decoder(nn.Module):
+  """Upscales quantized vectors to mask."""
+
+  @nn.compact
+  def __call__(self, x):
+    num_res_blocks = 2
+    dim = 128
+    num_upsample_layers = 4
+
+    x = nn.Conv(features=dim, kernel_size=(1, 1), padding=0)(x)
+    x = nn.relu(x)
+
+    for _ in range(num_res_blocks):
+      x = ResBlock(features=dim)(x)
+
+    for _ in range(num_upsample_layers):
+      x = nn.ConvTranspose(
+          features=dim,
+          kernel_size=(4, 4),
+          strides=(2, 2),
+          padding=2,
+          transpose_kernel=True,
+      )(x)
+      x = nn.relu(x)
+      dim //= 2
+
+    x = nn.Conv(features=1, kernel_size=(1, 1), padding=0)(x)
+
+    return x
+
+
+@functools.cache
+def get_reconstruct_masks(model):
+  """Reconstructs masks from codebook indices.
+
+  Based on code from https://arxiv.org/abs/2301.02229
+
+  Verified in
+  https://colab.research.google.com/drive/1AOr0cokOpM6-N9Z5HmxoeGxGj6jS37Vl
+
+  Args:
+    model: Model to use for conversion.
+
+  Returns:
+    A function that expects indices shaped `[B, 16]` of dtype int32, each
+    ranging from 0 to 127 (inclusive), and that returns a decoded masks sized
+    `[B, 64, 64, 1]`, of dtype float32, in range [-1, 1].
+  """
+  def reconstruct_masks(codebook_indices):
+    quantized = _quantized_values_from_codebook_indices(
+        codebook_indices, params['_embeddings']
+    )
+    return Decoder().apply({'params': params}, quantized)
+
+  with gfile.GFile(_KNOWN_MODELS.get(model, model), 'rb') as f:
+    params = _get_params(dict(np.load(f)))
+
+  return jax.jit(reconstruct_masks, backend='cpu')

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/storepreds.py

@@ -0,0 +1,77 @@
+# 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.
+
+"""Evaluator to run inference and store results."""
+import functools
+
+import big_vision.evaluators.common as c
+import big_vision.input_pipeline
+import big_vision.pp.builder
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+
+import jax
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+
+
+class Evaluator:
+  """Evaluator to run inference and store results."""
+
+  def __init__(
+      self, predict_fn, tokenizer=None,
+      preds_outfile="{workdir}/{name}_{split}_preds.json",
+      annot_outfile="{workdir}/{name}_{split}_annotations.json",
+      id_key="id",
+      *, data, devices, **kw
+  ):
+    self.id_key = id_key
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={id_key}, data=data, devices=devices, **kw)
+
+    self.preds_outfile = c.resolve_outfile(
+        preds_outfile, name=data.get("name"), split=data.get("split", ""))
+    self.annot_outfile = c.resolve_outfile(
+        annot_outfile, name=data.get("name"), split=data.get("split", ""))
+
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token)
+
+  def run(self, train_state):
+    """Run eval."""
+    res = []
+
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded generated tokens.
+      tokens = self.decode(train_state, batch)
+
+      # (local_batch,)
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      image_ids = batch[self.id_key][ex_masks]
+      pred_captions = self.tok.to_str(tokens[ex_masks])
+
+      for image_id, caption in zip(image_ids, pred_captions):
+        res.append({self.id_key: str(image_id), "caption": caption})
+
+    res = c.multiprocess_write_json(self.preds_outfile, res)
+
+    if jax.process_index():  # Host0 gets all preds and does eval.
+      return
+
+    yield "num_examples", len(res)

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/tallyqa.py

@@ -0,0 +1,144 @@
+# 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.
+
+"""Evaluator for TallyQA dataset."""
+
+import functools
+
+import big_vision.evaluators.common as c
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+
+
+# Largest count we want to track.
+_LARGEST_COUNT = 15
+
+
+class Evaluator:
+  """TallyQA evaluator."""
+
+  def __init__(self, predict_fn, tokenizer, *, devices, **kw):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={"answer", "issimple"}, devices=devices, **kw)
+
+    # We'll need the tokenizer to detokenize the model outputs later.
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token
+    )
+
+  def run(self, train_state):
+    """Does one evaluation run, yields metrics."""
+
+    accuracies_by_type = {"all": [], "simple": [], "complex": []}
+    # Add per-count entries. Cannot use a `defaultdict` as we need to `tree_map`
+    # over keys later in `c.process_sum`.
+    accuracies_by_type.update(
+        {f"count_{i}": [] for i in range(_LARGEST_COUNT + 1)}
+    )
+
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded (generated) token sequences suffixes.
+      tokens = self.decode(train_state, batch)
+
+      # (local_batch,) that indicates padding examples (0) vs real examples (1).
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      # Turn predictions into texts and then scores, one by one.
+      # We always compare the gt (string digit, e.g. "1") to the answer by the
+      # model (e.g. "1").
+      for i in range(len(tokens)):
+        if ex_masks[i] == 0:  # Skip last-batch padding examples
+          continue
+
+        # Extract the suffix/answer from the generated string, skip bos.
+        answer = self.tok.to_str(tokens[i], stop_at_eos=True)
+        # Standardize the reponse, i.e., convert number words ("one") to
+        # numerals ("1").
+        answer = _number_word_to_numeral(answer)
+
+        # Always need to do light space-processing:
+        gt = _number_word_to_numeral(batch["answer"][i])
+        accuracies_by_type["all"].append(float(answer == gt))
+
+        if "issimple" in batch:
+          # Simple/complex split.
+          if batch["issimple"][i] == 1:
+            accuracies_by_type["simple"].append(float(answer == gt))
+          elif batch["issimple"][i] == 0:
+            accuracies_by_type["complex"].append(float(answer == gt))
+          else:
+            # Train set is not annotated with simple/complex (but has dummy
+            # value of `-1` in this field).
+            pass
+
+        # Store accuracies per count.
+        accuracies_by_type[f"count_{gt}"].append(float(answer == gt))
+
+    # At this point `accuracies` is a list of per-example scores. However,
+    # remember that each host holds a different subset of the examples! So if
+    # we were to just return the mean accuracy here, we would effectively only
+    # have evaluated on the main host's (who writes metrics) subset!
+    # So now, we need to compute global means.
+    # There is one more caveat: `process_sum` needs the summands on each host
+    # to have the same size. So we either need to include dummy values for
+    # the padding examples (last batch, annoying), or we only sum scalars as in
+    # sufficient statistics, which we do here.
+    sum_accs = c.process_sum({k: sum(v) for k, v in accuracies_by_type.items()})
+    num_accs = c.process_sum({k: len(v) for k, v in accuracies_by_type.items()})
+
+    if n := num_accs["all"]:
+      yield "acc", sum_accs["all"] / n
+      yield "num", n  # Just for sanity checks.
+    for key in sum_accs.keys():
+      if (key != "all") and (num_accs[key]):
+        yield f"acc/{key}", sum_accs[key] / num_accs[key]
+        yield f"num/{key}", num_accs[key]  # Just for sanity checks.
+
+
+def _number_word_to_numeral(s: str) -> str:
+  """Returns numeral for a given number word, e.g., "one" -> "1" (up to 20)."""
+  return REPLACEMENTS.get(s.lower(), s)
+
+
+REPLACEMENTS = {
+    "none": "0",
+    "zero": "0",
+    "one": "1",
+    "two": "2",
+    "three": "3",
+    "four": "4",
+    "five": "5",
+    "six": "6",
+    "seven": "7",
+    "eight": "8",
+    "nine": "9",
+    "ten": "10",
+    "eleven": "11",
+    "twelve": "12",
+    "thirteen": "13",
+    "fourteen": "14",
+    "fifteen": "15",
+    "sixteen": "16",
+    "seventeen": "17",
+    "eighteen": "18",
+    "nineteen": "19",
+    "twenty": "20",
+}

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/vqa.py

@@ -0,0 +1,163 @@
+# 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.
+
+"""Evaluator for simple VQA variants (OCR-VQA, OKVQA, A-OKVQA).
+
+According to the (A-)OKVAQ papers, the eval for these datasets should follow
+VQAv2. But here we don't track different answer-types, and don't do any
+leave-one-out averaging, as this isn't done in the official implementation at
+https://github.com/allenai/aokvqa/blob/main/evaluation/eval_predictions.py
+either.
+
+Please read the description of how evaluators work at (internal link).
+This evaluator follows the pattern of also parallelizing the CPU computations
+(ie postprocessing, score computation) across hosts for more scalability.
+
+For now, simple decoding is implemented as part of the evaluator. We'll soon
+unify and move to a library of decoding functions, including fancier and more
+efficient ones.
+"""
+import functools
+
+import big_vision.evaluators.common as c
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+import editdistance
+
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+
+
+class Evaluator:
+  """Evaluator for simple VQA tasks.
+
+  This evaluator expects the batch to contain a field `question_id` and a field
+  `answer` for single ground truth or `answers` for multiple ground truths.
+
+  The field names used when writting the json result can be controlled with
+  `out_question_key` and `out_answer_key`.
+  """
+
+  def __init__(
+      self, predict_fn, tokenizer, to_lower=False,
+      outfile="{workdir}/{split}.json",
+      out_question_key="question_id", out_answer_key="answer",
+      *, data, devices, **kw):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={"answers", "answer", "question_id"},
+        data=data, devices=devices, **kw)
+
+    self.outfile = c.resolve_outfile(outfile, split=data.get("split"))
+    self.out_question_key = out_question_key
+    self.out_answer_key = out_answer_key
+
+    # We'll need the tokenizer to detokenize the model outputs later.
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.postproc = (lambda s: s.lower()) if to_lower else lambda s: s
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token)
+
+  def run(self, train_state):
+    """Does one evaluation run, yields metrics."""
+
+    accuracies = []
+    accuracies_any = []
+    anls_values = []
+    json_out = []
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded generated tokens.
+      tokens = self.decode(train_state, batch)
+
+      # (local_batch,) that indicates padding examples (0) vs real examples (1).
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      # Turn predictions into texts and then scores, one by one.
+      for i in range(len(tokens)):
+        if ex_masks[i] == 0:  # Skip last-batch padding examples
+          continue
+
+        answer = self.postproc(self.tok.to_str(tokens[i], stop_at_eos=True))
+
+        # Now we have two commonly used VQA evaluation modes:
+        if "answer" in batch:
+          # single GT (eg ocrvqa): just compare to that answer, done.
+          gt = self.postproc(batch["answer"][i])
+          gts = [gt]
+          accuracies.append(float(answer == gt))
+          accuracies_any.append(float(answer == gt))
+          anls_values.append(anls_metric(gt, answer))
+        elif "answers" in batch and (gt_answers := batch["answers"][i]).size:
+          # multiple GTs (eg okvqa): introduced by VQA, compare to each of them
+          # with a threshold, see also: https://visualqa.org/evaluation.html
+          gts = [self.postproc(a) for a in gt_answers]
+          num_match = sum([answer == gt for gt in gts])
+          accuracies.append(min(1.0, num_match / 3.0))
+          accuracies_any.append(min(1.0, float(num_match)))
+          anls_values.append(max(anls_metric(gt, answer) for gt in gts))
+        else:
+          gts = []
+
+        json_out.append({
+            self.out_question_key: batch["question_id"][i].item(),
+            self.out_answer_key: answer} | ({"gts": gts} if gts else {}))
+
+    # At this point `accuracies` is a list of per-example scores. However,
+    # remember that each host holds a different subset of the examples! So if
+    # we were to just return the mean accuracy here, we would effectively only
+    # have evaluated on the main host's (who writes metrics) subset!
+    # So now, we need to compute global means.
+    # There is one more caveat: `process_sum` needs the summands on each host
+    # to have the same size. So we either need to include dummy values for
+    # the padding examples (last batch, annoying), or we only sum scalars as in
+    # sufficient statistics, which we do here.
+    sum_accs, sum_accs_any, sum_anls, num_accs, num = c.process_sum(
+        [sum(accuracies), sum(accuracies_any), sum(anls_values),
+         len(accuracies), len(json_out)])
+
+    # Yielding metric_name, value means logging the metric.
+    if num_accs:
+      yield "acc", sum_accs / num_accs
+      yield "acc_any", sum_accs_any / num_accs
+      yield "anls", sum_anls / num_accs
+
+    yield "num", num  # Just for sanity checks.
+    c.multiprocess_write_json(self.outfile, json_out)
+
+
+def anls_metric(target: str, prediction: str, theta: float = 0.5):
+  """Calculates ANLS for DocVQA.
+
+  There does not seem to be an official evaluation script.
+  Public implementation on which this implementation is based:
+  https://github.com/herobd/layoutlmv2/blob/main/eval_docvqa.py#L92
+
+  Original paper (see Eq 1): https://arxiv.org/pdf/1907.00490.pdf
+
+  Args:
+    target: Target string.
+    prediction: Predicted string.
+    theta: Filter threshold set to 0.5 for DocVQA.
+
+  Returns:
+    ANLS score.
+  """
+  if target:
+    edit_distance = editdistance.eval(target, prediction)
+    normalized_ld = edit_distance / max(len(target), len(prediction))
+    return 1 - normalized_ld if normalized_ld < theta else 0
+  else:
+    return float(prediction == "")

Tipsomaly/model/big_vision/evaluators/proj/paligemma/transfers/vqav2.py

@@ -0,0 +1,197 @@
+# 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.
+
+"""Evaluator for VQAV2 dataset.
+"""
+import functools
+import re
+
+import big_vision.evaluators.common as c
+import big_vision.pp.tokenizer
+import big_vision.utils as u
+import numpy as np
+
+
+# Temporary global flag to facilitate backwards compatability. Will be removed
+# by the end of year 2023.
+API = "jit"
+
+
+class Evaluator:
+  """VQAv2 evaluator."""
+
+  def __init__(
+      self, predict_fn, tokenizer, outfile="{workdir}/{split}.json",
+      *, data, devices, **kw):
+    self.get_data_iter, self.steps = c.eval_input_pipeline(
+        keep_on_cpu={"answers", "answer_type", "question_type", "question_id"},
+        data=data, devices=devices, **kw)
+
+    self.outfile = c.resolve_outfile(outfile, split=data.get("split"))
+
+    # We'll need the tokenizer to detokenize the model outputs later.
+    self.tok = big_vision.pp.tokenizer.get_tokenizer(tokenizer)
+    self.decode = functools.partial(
+        predict_fn, devices=devices, eos_token=self.tok.eos_token)
+
+  def run(self, train_state):
+    """Does one evaluation run, yields metrics."""
+    accuracies_by_type = {"yes/no": [], "number": [], "other": []}
+    json_out = []
+
+    for _, batch in zip(range(self.steps), self.get_data_iter()):
+      # (batch, seqlen) array of decoded (generated) token sequences suffixes.
+      tokens = self.decode(train_state, batch)
+
+      # (local_batch,) that indicates padding examples (0) vs real examples (1).
+      tokens = u.get_local_slice_from_fsarray(tokens)
+      ex_masks = u.get_local_slice_from_fsarray(batch["_mask"])
+
+      # Turn predictions into texts and then scores, one by one.
+      for i in range(len(tokens)):
+        if ex_masks[i] == 0:  # Skip last-batch padding examples
+          continue
+
+        # Extract the suffix/answer from the generated string, skip bos.
+        answer = self.tok.to_str(tokens[i], stop_at_eos=True)
+        json = {"question_id": batch["question_id"][i].item(), "answer": answer}
+
+        # The rest is computation of VQA-score which compares to multiple GTs.
+        # This is described better here: https://visualqa.org/evaluation.html
+        if (gt_answers := batch["answers"][i]).size:
+          # Always need to do light space-processing:
+          gt_answers = [stripspace_vqav2(a) for a in gt_answers]
+          answer = stripspace_vqav2(answer)
+
+          # Only post-process if not all agree. Supposedly avoids postproc OCR:
+          # https://github.com/GT-Vision-Lab/VQA/issues/14#issuecomment-1334695361
+          if len(set(gt_answers)) > 1:
+            answer = postprocess_vqav2_text(answer)
+            gt_answers = [postprocess_vqav2_text(a) for a in gt_answers]
+
+          # Accuracy is avg over all ten leave-one-out GT's.
+          # https://github.com/GT-Vision-Lab/VQA/issues/1#issuecomment-199921352
+          # An answer is counted 100% correct as soon as 3 GT's agree with it.
+          matches = answer == np.array(gt_answers)
+          acc = np.mean([
+              np.clip(np.sum(np.delete(matches, i_leave_out)) / 3, 0, 1)
+              for i_leave_out in range(10)
+          ])
+
+          accuracies_by_type[batch["answer_type"][i]].append(acc)
+
+          # Update json with fully post-processed answer and gt:
+          json["answer_raw"] = json["answer"]
+          json["answer"] = answer
+          json["gts"] = gt_answers
+
+        json_out.append(json)
+
+    # At this point `accuracies` is a list of per-example scores. However,
+    # remember that each host holds a different subset of the examples! So if
+    # we were to just return the mean accuracy here, we would effectively only
+    # have evaluated on the main host's (who writes metrics) subset!
+    # So now, we need to compute global means.
+    # There is one more caveat: `process_sum` needs the summands on each host
+    # to have the same size. So we either need to include dummy values for
+    # the padding examples (last batch, annoying), or we only sum scalars as in
+    # sufficient statistics, which we do here.
+    sum_accs = c.process_sum({k: sum(v) for k, v in accuracies_by_type.items()})
+    num_accs = c.process_sum({k: len(v) for k, v in accuracies_by_type.items()})
+    num = c.process_sum(len(json_out))
+
+    # Yielding metric_name, value means logging the metric.
+    if n := sum(num_accs.values()):
+      yield "acc", sum(sum_accs.values()) / n
+    if n := num_accs["yes/no"]:
+      yield "acc/yesno", sum_accs["yes/no"] / n
+      yield "num/yesno", n
+    if n := num_accs["number"]:
+      yield "acc/number", sum_accs["number"] / n
+      yield "num/number", n
+    if n := num_accs["other"]:
+      yield "acc/other", sum_accs["other"] / n
+      yield "num/other", n
+
+    yield "num", num  # Just for sanity checks.
+    c.multiprocess_write_json(self.outfile, json_out)
+
+
+# Post-processing required is described at https://visualqa.org/evaluation.html
+
+
+def stripspace_vqav2(txt):
+  return txt.replace("\n", " ").replace("\t", " ").strip()
+
+
+def postprocess_vqav2_text(txt):
+  """Cleanup string according to VQA."""
+  has_digit_comma = re.search(r"(\d)(\,)(\d)", txt) is not None
+
+  out = txt
+  for p in PUNCT:
+    # NOTE: digit_comma here looks like a bug in official code, so we follow it.
+    if has_digit_comma or f"{p} " in txt or f" {p}" in txt:
+      out = out.replace(p, "")
+    else:
+      out = out.replace(p, " ")
+
+  # Remove full-stops that aren't part of a number.
+  out = re.sub(r"(?!<=\d)(\.)(?!\d)", "", out, flags=re.UNICODE)
+
+  words = []
+  for word in out.lower().split():
+    if word not in ARTICLES:
+      words.append(REPLACEMENTS.get(word, word))
+  return " ".join(words)
+
+
+# pylint: disable=line-too-long
+REPLACEMENTS = {
+    # CONTRACTIONS
+    "aint": "ain't", "arent": "aren't", "cant": "can't", "couldve": "could've", "couldnt": "couldn't",
+    "couldn'tve": "couldn't've", "couldnt've": "couldn't've", "didnt": "didn't", "doesnt": "doesn't", "dont": "don't", "hadnt": "hadn't",
+    "hadnt've": "hadn't've", "hadn'tve": "hadn't've", "hasnt": "hasn't", "havent": "haven't", "hed": "he'd", "hed've": "he'd've",
+    "he'dve": "he'd've", "hes": "he's", "howd": "how'd", "howll": "how'll", "hows": "how's", "Id've": "I'd've", "I'dve": "I'd've",
+    "Im": "I'm", "Ive": "I've", "isnt": "isn't", "itd": "it'd", "itd've": "it'd've", "it'dve": "it'd've", "itll": "it'll", "let's": "let's",
+    "maam": "ma'am", "mightnt": "mightn't", "mightnt've": "mightn't've", "mightn'tve": "mightn't've", "mightve": "might've",
+    "mustnt": "mustn't", "mustve": "must've", "neednt": "needn't", "notve": "not've", "oclock": "o'clock", "oughtnt": "oughtn't",
+    "ow's'at": "'ow's'at", "'ows'at": "'ow's'at", "'ow'sat": "'ow's'at", "shant": "shan't", "shed've": "she'd've", "she'dve": "she'd've",
+    "she's": "she's", "shouldve": "should've", "shouldnt": "shouldn't", "shouldnt've": "shouldn't've", "shouldn'tve": "shouldn't've",
+    "somebody'd": "somebodyd", "somebodyd've": "somebody'd've", "somebody'dve": "somebody'd've", "somebodyll": "somebody'll",
+    "somebodys": "somebody's", "someoned": "someone'd", "someoned've": "someone'd've", "someone'dve": "someone'd've",
+    "someonell": "someone'll", "someones": "someone's", "somethingd": "something'd", "somethingd've": "something'd've",
+    "something'dve": "something'd've", "somethingll": "something'll", "thats": "that's", "thered": "there'd", "thered've": "there'd've",
+    "there'dve": "there'd've", "therere": "there're", "theres": "there's", "theyd": "they'd", "theyd've": "they'd've",
+    "they'dve": "they'd've", "theyll": "they'll", "theyre": "they're", "theyve": "they've", "twas": "'twas", "wasnt": "wasn't",
+    "wed've": "we'd've", "we'dve": "we'd've", "weve": "we've", "werent": "weren't", "whatll": "what'll", "whatre": "what're",
+    "whats": "what's", "whatve": "what've", "whens": "when's", "whered": "where'd", "wheres": "where's", "whereve": "where've",
+    "whod": "who'd", "whod've": "who'd've", "who'dve": "who'd've", "wholl": "who'll", "whos": "who's", "whove": "who've", "whyll": "why'll",
+    "whyre": "why're", "whys": "why's", "wont": "won't", "wouldve": "would've", "wouldnt": "wouldn't", "wouldnt've": "wouldn't've",
+    "wouldn'tve": "wouldn't've", "yall": "y'all", "yall'll": "y'all'll", "y'allll": "y'all'll", "yall'd've": "y'all'd've",
+    "y'alld've": "y'all'd've", "y'all'dve": "y'all'd've", "youd": "you'd", "youd've": "you'd've", "you'dve": "you'd've",
+    "youll": "you'll", "youre": "you're", "youve": "you've",
+    # NUMBERS
+    "none": "0", "zero": "0", "one": "1", "two": "2",
+    "three": "3", "four": "4", "five": "5", "six": "6",
+    "seven": "7", "eight": "8", "nine": "9", "ten": "10",
+}
+# pylint: enable=line-too-long
+
+PUNCT = [
+    ";", "/", "[", "]", "\"", "{", "}",
+    "(", ")", "=", "+", "\\", "_", "-",
+    ">", "<", "@", "`", ",", "?", "!"
+]
+ARTICLES = {"a", "an", "the"}

Tipsomaly/model/big_vision/pp/archive/autoaugment.py

@@ -0,0 +1,700 @@
+# 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.
+
+"""AutoAugment and RandAugment policies for enhanced image preprocessing.
+
+AutoAugment Reference: https://arxiv.org/abs/1805.09501
+RandAugment Reference: https://arxiv.org/abs/1909.13719
+
+This code is forked from
+https://github.com/tensorflow/tpu/blob/11d0db15cf1c3667f6e36fecffa111399e008acd/models/official/efficientnet/autoaugment.py
+"""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import dataclasses
+import inspect
+import math
+import tensorflow.compat.v1 as tf
+from tensorflow_addons import image as contrib_image
+
+# This signifies the max integer that the controller RNN could predict for the
+# augmentation scheme.
+_MAX_LEVEL = 10.
+
+
+@dataclasses.dataclass
+class HParams:
+  """Parameters for AutoAugment and RandAugment."""
+  cutout_const: int
+  translate_const: int
+
+
+def policy_v0():
+  """Autoaugment policy that was used in AutoAugment Paper."""
+  # Each tuple is an augmentation operation of the form
+  # (operation, probability, magnitude). Each element in policy is a
+  # sub-policy that will be applied sequentially on the image.
+  policy = [
+      [('Equalize', 0.8, 1), ('ShearY', 0.8, 4)],
+      [('Color', 0.4, 9), ('Equalize', 0.6, 3)],
+      [('Color', 0.4, 1), ('Rotate', 0.6, 8)],
+      [('Solarize', 0.8, 3), ('Equalize', 0.4, 7)],
+      [('Solarize', 0.4, 2), ('Solarize', 0.6, 2)],
+      [('Color', 0.2, 0), ('Equalize', 0.8, 8)],
+      [('Equalize', 0.4, 8), ('SolarizeAdd', 0.8, 3)],
+      [('ShearX', 0.2, 9), ('Rotate', 0.6, 8)],
+      [('Color', 0.6, 1), ('Equalize', 1.0, 2)],
+      [('Invert', 0.4, 9), ('Rotate', 0.6, 0)],
+      [('Equalize', 1.0, 9), ('ShearY', 0.6, 3)],
+      [('Color', 0.4, 7), ('Equalize', 0.6, 0)],
+      [('Posterize', 0.4, 6), ('AutoContrast', 0.4, 7)],
+      [('Solarize', 0.6, 8), ('Color', 0.6, 9)],
+      [('Solarize', 0.2, 4), ('Rotate', 0.8, 9)],
+      [('Rotate', 1.0, 7), ('TranslateY', 0.8, 9)],
+      [('ShearX', 0.0, 0), ('Solarize', 0.8, 4)],
+      [('ShearY', 0.8, 0), ('Color', 0.6, 4)],
+      [('Color', 1.0, 0), ('Rotate', 0.6, 2)],
+      [('Equalize', 0.8, 4), ('Equalize', 0.0, 8)],
+      [('Equalize', 1.0, 4), ('AutoContrast', 0.6, 2)],
+      [('ShearY', 0.4, 7), ('SolarizeAdd', 0.6, 7)],
+      [('Posterize', 0.8, 2), ('Solarize', 0.6, 10)],
+      [('Solarize', 0.6, 8), ('Equalize', 0.6, 1)],
+      [('Color', 0.8, 6), ('Rotate', 0.4, 5)],
+  ]
+  return policy
+
+
+def policy_vtest():
+  """Autoaugment test policy for debugging."""
+  # Each tuple is an augmentation operation of the form
+  # (operation, probability, magnitude). Each element in policy is a
+  # sub-policy that will be applied sequentially on the image.
+  policy = [
+      [('TranslateX', 1.0, 4), ('Equalize', 1.0, 10)],
+  ]
+  return policy
+
+
+def blend(image1, image2, factor):
+  """Blend image1 and image2 using 'factor'.
+  Factor can be above 0.0.  A value of 0.0 means only image1 is used.
+  A value of 1.0 means only image2 is used.  A value between 0.0 and
+  1.0 means we linearly interpolate the pixel values between the two
+  images.  A value greater than 1.0 "extrapolates" the difference
+  between the two pixel values, and we clip the results to values
+  between 0 and 255.
+  Args:
+    image1: An image Tensor of type uint8.
+    image2: An image Tensor of type uint8.
+    factor: A floating point value above 0.0.
+  Returns:
+    A blended image Tensor of type uint8.
+  """
+  if factor == 0.0:
+    return tf.convert_to_tensor(image1)
+  if factor == 1.0:
+    return tf.convert_to_tensor(image2)
+
+  image1 = tf.to_float(image1)
+  image2 = tf.to_float(image2)
+
+  difference = image2 - image1
+  scaled = factor * difference
+
+  # Do addition in float.
+  temp = tf.to_float(image1) + scaled
+
+  # Interpolate
+  if factor > 0.0 and factor < 1.0:
+    # Interpolation means we always stay within 0 and 255.
+    return tf.cast(temp, tf.uint8)
+
+  # Extrapolate:
+  #
+  # We need to clip and then cast.
+  return tf.cast(tf.clip_by_value(temp, 0.0, 255.0), tf.uint8)
+
+
+def cutout(image, pad_size, replace=0):
+  """Apply cutout (https://arxiv.org/abs/1708.04552) to image.
+  This operation applies a (2*pad_size x 2*pad_size) mask of zeros to
+  a random location within `img`. The pixel values filled in will be of the
+  value `replace`. The located where the mask will be applied is randomly
+  chosen uniformly over the whole image.
+  Args:
+    image: An image Tensor of type uint8.
+    pad_size: Specifies how big the zero mask that will be generated is that
+      is applied to the image. The mask will be of size
+      (2*pad_size x 2*pad_size).
+    replace: What pixel value to fill in the image in the area that has
+      the cutout mask applied to it.
+  Returns:
+    An image Tensor that is of type uint8.
+  """
+  image_height = tf.shape(image)[0]
+  image_width = tf.shape(image)[1]
+
+  # Sample the center location in the image where the zero mask will be applied.
+  cutout_center_height = tf.random_uniform(
+      shape=[], minval=0, maxval=image_height,
+      dtype=tf.int32)
+
+  cutout_center_width = tf.random_uniform(
+      shape=[], minval=0, maxval=image_width,
+      dtype=tf.int32)
+
+  lower_pad = tf.maximum(0, cutout_center_height - pad_size)
+  upper_pad = tf.maximum(0, image_height - cutout_center_height - pad_size)
+  left_pad = tf.maximum(0, cutout_center_width - pad_size)
+  right_pad = tf.maximum(0, image_width - cutout_center_width - pad_size)
+
+  cutout_shape = [image_height - (lower_pad + upper_pad),
+                  image_width - (left_pad + right_pad)]
+  padding_dims = [[lower_pad, upper_pad], [left_pad, right_pad]]
+  mask = tf.pad(
+      tf.zeros(cutout_shape, dtype=image.dtype),
+      padding_dims, constant_values=1)
+  mask = tf.expand_dims(mask, -1)
+  mask = tf.tile(mask, [1, 1, 3])
+  image = tf.where(
+      tf.equal(mask, 0),
+      tf.ones_like(image, dtype=image.dtype) * replace,
+      image)
+  return image
+
+
+def solarize(image, threshold=128):
+  # For each pixel in the image, select the pixel
+  # if the value is less than the threshold.
+  # Otherwise, subtract 255 from the pixel.
+  return tf.where(image < threshold, image, 255 - image)
+
+
+def solarize_add(image, addition=0, threshold=128):
+  # For each pixel in the image less than threshold
+  # we add 'addition' amount to it and then clip the
+  # pixel value to be between 0 and 255. The value
+  # of 'addition' is between -128 and 128.
+  added_image = tf.cast(image, tf.int64) + addition
+  added_image = tf.cast(tf.clip_by_value(added_image, 0, 255), tf.uint8)
+  return tf.where(image < threshold, added_image, image)
+
+
+def color(image, factor):
+  """Equivalent of PIL Color."""
+  degenerate = tf.image.grayscale_to_rgb(tf.image.rgb_to_grayscale(image))
+  return blend(degenerate, image, factor)
+
+
+def contrast(image, factor):
+  """Equivalent of PIL Contrast."""
+  degenerate = tf.image.rgb_to_grayscale(image)
+  # Cast before calling tf.histogram.
+  degenerate = tf.cast(degenerate, tf.int32)
+
+  # Compute the grayscale histogram, then compute the mean pixel value,
+  # and create a constant image size of that value.  Use that as the
+  # blending degenerate target of the original image.
+  hist = tf.histogram_fixed_width(degenerate, [0, 255], nbins=256)
+  mean = tf.reduce_sum(tf.cast(hist, tf.float32)) / 256.0
+  degenerate = tf.ones_like(degenerate, dtype=tf.float32) * mean
+  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
+  degenerate = tf.image.grayscale_to_rgb(tf.cast(degenerate, tf.uint8))
+  return blend(degenerate, image, factor)
+
+
+def brightness(image, factor):
+  """Equivalent of PIL Brightness."""
+  degenerate = tf.zeros_like(image)
+  return blend(degenerate, image, factor)
+
+
+def posterize(image, bits):
+  """Equivalent of PIL Posterize."""
+  shift = 8 - bits
+  return tf.bitwise.left_shift(tf.bitwise.right_shift(image, shift), shift)
+
+
+def rotate(image, degrees, replace):
+  """Rotates the image by degrees either clockwise or counterclockwise.
+  Args:
+    image: An image Tensor of type uint8.
+    degrees: Float, a scalar angle in degrees to rotate all images by. If
+      degrees is positive the image will be rotated clockwise otherwise it will
+      be rotated counterclockwise.
+    replace: A one or three value 1D tensor to fill empty pixels caused by
+      the rotate operation.
+  Returns:
+    The rotated version of image.
+  """
+  # Convert from degrees to radians.
+  degrees_to_radians = math.pi / 180.0
+  radians = degrees * degrees_to_radians
+
+  # In practice, we should randomize the rotation degrees by flipping
+  # it negatively half the time, but that's done on 'degrees' outside
+  # of the function.
+  image = contrib_image.rotate(wrap(image), radians)
+  return unwrap(image, replace)
+
+
+def translate_x(image, pixels, replace):
+  """Equivalent of PIL Translate in X dimension."""
+  image = contrib_image.translate(wrap(image), [-pixels, 0])
+  return unwrap(image, replace)
+
+
+def translate_y(image, pixels, replace):
+  """Equivalent of PIL Translate in Y dimension."""
+  image = contrib_image.translate(wrap(image), [0, -pixels])
+  return unwrap(image, replace)
+
+
+def shear_x(image, level, replace):
+  """Equivalent of PIL Shearing in X dimension."""
+  # Shear parallel to x axis is a projective transform
+  # with a matrix form of:
+  # [1  level
+  #  0  1].
+  image = contrib_image.transform(
+      wrap(image), [1., level, 0., 0., 1., 0., 0., 0.])
+  return unwrap(image, replace)
+
+
+def shear_y(image, level, replace):
+  """Equivalent of PIL Shearing in Y dimension."""
+  # Shear parallel to y axis is a projective transform
+  # with a matrix form of:
+  # [1  0
+  #  level  1].
+  image = contrib_image.transform(
+      wrap(image), [1., 0., 0., level, 1., 0., 0., 0.])
+  return unwrap(image, replace)
+
+
+def autocontrast(image):
+  """Implements Autocontrast function from PIL using TF ops.
+  Args:
+    image: A 3D uint8 tensor.
+  Returns:
+    The image after it has had autocontrast applied to it and will be of type
+    uint8.
+  """
+
+  def scale_channel(image):
+    """Scale the 2D image using the autocontrast rule."""
+    # A possibly cheaper version can be done using cumsum/unique_with_counts
+    # over the histogram values, rather than iterating over the entire image.
+    # to compute mins and maxes.
+    lo = tf.to_float(tf.reduce_min(image))
+    hi = tf.to_float(tf.reduce_max(image))
+
+    # Scale the image, making the lowest value 0 and the highest value 255.
+    def scale_values(im):
+      scale = 255.0 / (hi - lo)
+      offset = -lo * scale
+      im = tf.to_float(im) * scale + offset
+      im = tf.clip_by_value(im, 0.0, 255.0)
+      return tf.cast(im, tf.uint8)
+
+    result = tf.cond(hi > lo, lambda: scale_values(image), lambda: image)
+    return result
+
+  # Assumes RGB for now.  Scales each channel independently
+  # and then stacks the result.
+  s1 = scale_channel(image[:, :, 0])
+  s2 = scale_channel(image[:, :, 1])
+  s3 = scale_channel(image[:, :, 2])
+  image = tf.stack([s1, s2, s3], 2)
+  return image
+
+
+def sharpness(image, factor):
+  """Implements Sharpness function from PIL using TF ops."""
+  orig_image = image
+  image = tf.cast(image, tf.float32)
+  # Make image 4D for conv operation.
+  image = tf.expand_dims(image, 0)
+  # SMOOTH PIL Kernel.
+  kernel = tf.constant(
+      [[1, 1, 1], [1, 5, 1], [1, 1, 1]], dtype=tf.float32,
+      shape=[3, 3, 1, 1]) / 13.
+  # Tile across channel dimension.
+  kernel = tf.tile(kernel, [1, 1, 3, 1])
+  strides = [1, 1, 1, 1]
+  with tf.device('/cpu:0'):
+    # Some augmentation that uses depth-wise conv will cause crashing when
+    # training on GPU. See ((internal link)) for details.
+    degenerate = tf.nn.depthwise_conv2d(
+        image, kernel, strides, padding='VALID', rate=[1, 1])
+  degenerate = tf.clip_by_value(degenerate, 0.0, 255.0)
+  degenerate = tf.squeeze(tf.cast(degenerate, tf.uint8), [0])
+
+  # For the borders of the resulting image, fill in the values of the
+  # original image.
+  mask = tf.ones_like(degenerate)
+  padded_mask = tf.pad(mask, [[1, 1], [1, 1], [0, 0]])
+  padded_degenerate = tf.pad(degenerate, [[1, 1], [1, 1], [0, 0]])
+  result = tf.where(tf.equal(padded_mask, 1), padded_degenerate, orig_image)
+
+  # Blend the final result.
+  return blend(result, orig_image, factor)
+
+
+def equalize(image):
+  """Implements Equalize function from PIL using TF ops."""
+  def scale_channel(im, c):
+    """Scale the data in the channel to implement equalize."""
+    im = tf.cast(im[:, :, c], tf.int32)
+    # Compute the histogram of the image channel.
+    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)
+
+    # For the purposes of computing the step, filter out the nonzeros.
+    nonzero = tf.where(tf.not_equal(histo, 0))
+    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
+    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255
+
+    def build_lut(histo, step):
+      # Compute the cumulative sum, shifting by step // 2
+      # and then normalization by step.
+      lut = (tf.cumsum(histo) + (step // 2)) // step
+      # Shift lut, prepending with 0.
+      lut = tf.concat([[0], lut[:-1]], 0)
+      # Clip the counts to be in range.  This is done
+      # in the C code for image.point.
+      return tf.clip_by_value(lut, 0, 255)
+
+    # If step is zero, return the original image.  Otherwise, build
+    # lut from the full histogram and step and then index from it.
+    result = tf.cond(tf.equal(step, 0),
+                     lambda: im,
+                     lambda: tf.gather(build_lut(histo, step), im))
+
+    return tf.cast(result, tf.uint8)
+
+  # Assumes RGB for now.  Scales each channel independently
+  # and then stacks the result.
+  s1 = scale_channel(image, 0)
+  s2 = scale_channel(image, 1)
+  s3 = scale_channel(image, 2)
+  image = tf.stack([s1, s2, s3], 2)
+  return image
+
+
+def invert(image):
+  """Inverts the image pixels."""
+  image = tf.convert_to_tensor(image)
+  return 255 - image
+
+
+def wrap(image):
+  """Returns 'image' with an extra channel set to all 1s."""
+  shape = tf.shape(image)
+  extended_channel = tf.ones([shape[0], shape[1], 1], image.dtype)
+  extended = tf.concat([image, extended_channel], 2)
+  return extended
+
+
+def unwrap(image, replace):
+  """Unwraps an image produced by wrap.
+  Where there is a 0 in the last channel for every spatial position,
+  the rest of the three channels in that spatial dimension are grayed
+  (set to 128).  Operations like translate and shear on a wrapped
+  Tensor will leave 0s in empty locations.  Some transformations look
+  at the intensity of values to do preprocessing, and we want these
+  empty pixels to assume the 'average' value, rather than pure black.
+  Args:
+    image: A 3D Image Tensor with 4 channels.
+    replace: A one or three value 1D tensor to fill empty pixels.
+  Returns:
+    image: A 3D image Tensor with 3 channels.
+  """
+  image_shape = tf.shape(image)
+  # Flatten the spatial dimensions.
+  flattened_image = tf.reshape(image, [-1, image_shape[2]])
+
+  # Find all pixels where the last channel is zero.
+  alpha_channel = flattened_image[:, 3]
+
+  replace = tf.concat([replace, tf.ones([1], image.dtype)], 0)
+
+  # Where they are zero, fill them in with 'replace'.
+  flattened_image = tf.where(
+      tf.equal(alpha_channel, 0),
+      tf.ones_like(flattened_image, dtype=image.dtype) * replace,
+      flattened_image)
+
+  image = tf.reshape(flattened_image, image_shape)
+  image = tf.slice(image, [0, 0, 0], [image_shape[0], image_shape[1], 3])
+  return image
+
+
+NAME_TO_FUNC = {
+    'AutoContrast': autocontrast,
+    'Equalize': equalize,
+    'Invert': invert,
+    'Rotate': rotate,
+    'Posterize': posterize,
+    'Solarize': solarize,
+    'SolarizeAdd': solarize_add,
+    'Color': color,
+    'Contrast': contrast,
+    'Brightness': brightness,
+    'Sharpness': sharpness,
+    'ShearX': shear_x,
+    'ShearY': shear_y,
+    'TranslateX': translate_x,
+    'TranslateY': translate_y,
+    'Cutout': cutout,
+}
+
+
+def _randomly_negate_tensor(tensor):
+  """With 50% prob turn the tensor negative."""
+  should_flip = tf.cast(tf.floor(tf.random_uniform([]) + 0.5), tf.bool)
+  final_tensor = tf.cond(should_flip, lambda: tensor, lambda: -tensor)
+  return final_tensor
+
+
+def _rotate_level_to_arg(level):
+  level = (level/_MAX_LEVEL) * 30.
+  level = _randomly_negate_tensor(level)
+  return (level,)
+
+
+def _shrink_level_to_arg(level):
+  """Converts level to ratio by which we shrink the image content."""
+  if level == 0:
+    return (1.0,)  # if level is zero, do not shrink the image
+  # Maximum shrinking ratio is 2.9.
+  level = 2. / (_MAX_LEVEL / level) + 0.9
+  return (level,)
+
+
+def _enhance_level_to_arg(level):
+  return ((level/_MAX_LEVEL) * 1.8 + 0.1,)
+
+
+def _shear_level_to_arg(level):
+  level = (level/_MAX_LEVEL) * 0.3
+  # Flip level to negative with 50% chance.
+  level = _randomly_negate_tensor(level)
+  return (level,)
+
+
+def _translate_level_to_arg(level, translate_const):
+  level = (level/_MAX_LEVEL) * float(translate_const)
+  # Flip level to negative with 50% chance.
+  level = _randomly_negate_tensor(level)
+  return (level,)
+
+
+def level_to_arg(hparams):
+  return {
+      'AutoContrast': lambda level: (),
+      'Equalize': lambda level: (),
+      'Invert': lambda level: (),
+      'Rotate': _rotate_level_to_arg,
+      'Posterize': lambda level: (int((level/_MAX_LEVEL) * 4),),
+      'Solarize': lambda level: (int((level/_MAX_LEVEL) * 256),),
+      'SolarizeAdd': lambda level: (int((level/_MAX_LEVEL) * 110),),
+      'Color': _enhance_level_to_arg,
+      'Contrast': _enhance_level_to_arg,
+      'Brightness': _enhance_level_to_arg,
+      'Sharpness': _enhance_level_to_arg,
+      'ShearX': _shear_level_to_arg,
+      'ShearY': _shear_level_to_arg,
+      'Cutout': lambda level: (int((level/_MAX_LEVEL) * hparams.cutout_const),),
+      'TranslateX': lambda level: _translate_level_to_arg(
+          level, hparams.translate_const),
+      'TranslateY': lambda level: _translate_level_to_arg(
+          level, hparams.translate_const),
+      # pylint:enable=g-long-lambda
+  }
+
+
+def _parse_policy_info(name, prob, level, replace_value, augmentation_hparams):
+  """Return the function that corresponds to `name` and update `level` param."""
+  func = NAME_TO_FUNC[name]
+  args = level_to_arg(augmentation_hparams)[name](level)
+
+  # Check to see if prob is passed into function. This is used for operations
+  # where we alter bboxes independently.
+  # pytype:disable=wrong-arg-types
+  if 'prob' in inspect.getfullargspec(func).args:
+    args = tuple([prob] + list(args))
+  # pytype:enable=wrong-arg-types
+
+  # Add in replace arg if it is required for the function that is being called.
+  # pytype:disable=wrong-arg-types
+  if 'replace' in inspect.getfullargspec(func).args:
+    # Make sure replace is the final argument
+    assert 'replace' == inspect.getfullargspec(func).args[-1]
+    args = tuple(list(args) + [replace_value])
+  # pytype:enable=wrong-arg-types
+
+  return (func, prob, args)
+
+
+def _apply_func_with_prob(func, image, args, prob):
+  """Apply `func` to image w/ `args` as input with probability `prob`."""
+  assert isinstance(args, tuple)
+
+  # If prob is a function argument, then this randomness is being handled
+  # inside the function, so make sure it is always called.
+  # pytype:disable=wrong-arg-types
+  if 'prob' in inspect.getfullargspec(func).args:
+    prob = 1.0
+  # pytype:enable=wrong-arg-types
+
+  # Apply the function with probability `prob`.
+  should_apply_op = tf.cast(
+      tf.floor(tf.random_uniform([], dtype=tf.float32) + prob), tf.bool)
+  augmented_image = tf.cond(
+      should_apply_op,
+      lambda: func(image, *args),
+      lambda: image)
+  return augmented_image
+
+
+def select_and_apply_random_policy(policies, image):
+  """Select a random policy from `policies` and apply it to `image`."""
+  policy_to_select = tf.random_uniform([], maxval=len(policies), dtype=tf.int32)
+  # Note that using tf.case instead of tf.conds would result in significantly
+  # larger graphs and would even break export for some larger policies.
+  for (i, policy) in enumerate(policies):
+    image = tf.cond(
+        tf.equal(i, policy_to_select),
+        lambda selected_policy=policy: selected_policy(image),
+        lambda: image)
+  return image
+
+
+def build_and_apply_nas_policy(policies, image,
+                               augmentation_hparams):
+  """Build a policy from the given policies passed in and apply to image.
+  Args:
+    policies: list of lists of tuples in the form `(func, prob, level)`, `func`
+      is a string name of the augmentation function, `prob` is the probability
+      of applying the `func` operation, `level` is the input argument for
+      `func`.
+    image: tf.Tensor that the resulting policy will be applied to.
+    augmentation_hparams: Hparams associated with the NAS learned policy.
+  Returns:
+    A version of image that now has data augmentation applied to it based on
+    the `policies` pass into the function.
+  """
+  replace_value = [128, 128, 128]
+
+  # func is the string name of the augmentation function, prob is the
+  # probability of applying the operation and level is the parameter associated
+  # with the tf op.
+
+  # tf_policies are functions that take in an image and return an augmented
+  # image.
+  tf_policies = []
+  for policy in policies:
+    tf_policy = []
+    # Link string name to the correct python function and make sure the correct
+    # argument is passed into that function.
+    for policy_info in policy:
+      policy_info = list(policy_info) + [replace_value, augmentation_hparams]
+
+      tf_policy.append(_parse_policy_info(*policy_info))
+    # Now build the tf policy that will apply the augmentation procedue
+    # on image.
+    def make_final_policy(tf_policy_):
+      def final_policy(image_):
+        for func, prob, args in tf_policy_:
+          image_ = _apply_func_with_prob(
+              func, image_, args, prob)
+        return image_
+      return final_policy
+    tf_policies.append(make_final_policy(tf_policy))
+
+  augmented_image = select_and_apply_random_policy(
+      tf_policies, image)
+  return augmented_image
+
+
+def distort_image_with_autoaugment(image, augmentation_name):
+  """Applies the AutoAugment policy to `image`.
+  AutoAugment is from the paper: https://arxiv.org/abs/1805.09501.
+  Args:
+    image: `Tensor` of shape [height, width, 3] representing an image.
+    augmentation_name: The name of the AutoAugment policy to use. The available
+      options are `v0` and `test`. `v0` is the policy used for
+      all of the results in the paper and was found to achieve the best results
+      on the COCO dataset. `v1`, `v2` and `v3` are additional good policies
+      found on the COCO dataset that have slight variation in what operations
+      were used during the search procedure along with how many operations are
+      applied in parallel to a single image (2 vs 3).
+  Returns:
+    A tuple containing the augmented versions of `image`.
+  """
+  available_policies = {'v0': policy_v0,
+                        'test': policy_vtest}
+  if augmentation_name not in available_policies:
+    raise ValueError('Invalid augmentation_name: {}'.format(augmentation_name))
+
+  policy = available_policies[augmentation_name]()
+  # Hparams that will be used for AutoAugment.
+  augmentation_hparams = HParams(
+      cutout_const=100, translate_const=250)
+
+  return build_and_apply_nas_policy(policy, image, augmentation_hparams)
+
+
+def distort_image_with_randaugment(image, num_layers, magnitude):
+  """Applies the RandAugment policy to `image`.
+  RandAugment is from the paper https://arxiv.org/abs/1909.13719,
+  Args:
+    image: `Tensor` of shape [height, width, 3] representing an image.
+    num_layers: Integer, the number of augmentation transformations to apply
+      sequentially to an image. Represented as (N) in the paper. Usually best
+      values will be in the range [1, 3].
+    magnitude: Integer, shared magnitude across all augmentation operations.
+      Represented as (M) in the paper. Usually best values are in the range
+      [5, 30].
+  Returns:
+    The augmented version of `image`.
+  """
+  replace_value = [128] * 3
+  tf.logging.info('Using RandAug.')
+  augmentation_hparams = HParams(
+      cutout_const=40, translate_const=100)
+  available_ops = [
+      'AutoContrast', 'Equalize', 'Invert', 'Rotate', 'Posterize',
+      'Solarize', 'Color', 'Contrast', 'Brightness', 'Sharpness',
+      'ShearX', 'ShearY', 'TranslateX', 'TranslateY', 'Cutout', 'SolarizeAdd']
+
+  for layer_num in range(num_layers):
+    op_to_select = tf.random_uniform(
+        [], maxval=len(available_ops), dtype=tf.int32)
+    random_magnitude = float(magnitude)
+    with tf.name_scope('randaug_layer_{}'.format(layer_num)):
+      for (i, op_name) in enumerate(available_ops):
+        prob = tf.random_uniform([], minval=0.2, maxval=0.8, dtype=tf.float32)
+        func, _, args = _parse_policy_info(op_name, prob, random_magnitude,
+                                           replace_value, augmentation_hparams)
+        image = tf.cond(
+            tf.equal(i, op_to_select),
+            lambda selected_func=func, selected_args=args: selected_func(
+                image, *selected_args),
+            # pylint:enable=g-long-lambda
+            lambda: image)
+  return image

Tipsomaly/model/big_vision/pp/archive/randaug.py

@@ -0,0 +1,46 @@
+# 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.
+
+"""RandAug depends on deprecated tfa.image package, now defunct."""
+
+from big_vision.pp import registry
+from big_vision.pp import utils
+from big_vision.pp.archive import autoaugment
+
+
+@registry.Registry.register("preprocess_ops.randaug")
+@utils.InKeyOutKey()
+def get_randaug(num_layers: int = 2, magnitude: int = 10):
+  """Creates a function that applies RandAugment.
+
+  RandAugment is from the paper https://arxiv.org/abs/1909.13719,
+
+  Args:
+    num_layers: Integer, the number of augmentation transformations to apply
+      sequentially to an image. Represented as (N) in the paper. Usually best
+      values will be in the range [1, 3].
+    magnitude: Integer, shared magnitude across all augmentation operations.
+      Represented as (M) in the paper. Usually best values are in the range [5,
+      30].
+
+  Returns:
+    a function that applies RandAugment.
+  """
+
+  def _randaug(image):
+    return autoaugment.distort_image_with_randaugment(
+        image, num_layers, magnitude
+    )
+
+  return _randaug

Tipsomaly/model/big_vision/pp/proj/clippo/download_unifont.sh

@@ -0,0 +1,21 @@
+# Copyright 2022 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.
+
+#!/bin/bash
+# This is intended to be run from the big_vision repository root:
+#
+# bash big_vision/pp/proj/clippo/download_unifont.sh
+wget https://unifoundry.com/pub/unifont/unifont-9.0.06/font-builds/unifont-9.0.06.hex.gz https://unifoundry.com/pub/unifont/unifont-9.0.06/font-builds/unifont_upper-9.0.06.hex.gz
+gunzip unifont-9.0.06.hex.gz unifont_upper-9.0.06.hex.gz
+mv unifont-9.0.06.hex unifont_upper-9.0.06.hex big_vision/pp/proj/clippo/

Tipsomaly/model/big_vision/pp/proj/givt/pp_ops.py

@@ -0,0 +1,36 @@
+# 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.
+
+"""GIVT-specific preprocessing ops."""
+
+from big_vision.pp import registry
+from big_vision.pp import utils
+import tensorflow as tf
+
+
+@registry.Registry.register("preprocess_ops.bin_nyu_depth")
+@utils.InKeyOutKey(indefault="labels", outdefault="labels")
+def get_bin_nyu_depth(min_depth=0.001, max_depth=10.0, num_bins=256):
+  """Binning of NYU depth for UViM in preprocessing rather than model."""
+
+  def _bin_depth(labels):  # pylint: disable=missing-docstring
+    labels = (labels - min_depth) / (max_depth - min_depth)
+    labels *= num_bins
+    labels = tf.cast(tf.floor(labels), tf.int32)
+    labels = tf.minimum(labels, num_bins - 1)
+    labels = tf.maximum(labels, 0)
+    return labels
+
+  return _bin_depth
+

Tipsomaly/model/big_vision/pp/proj/image_text/ops_naflex.py

@@ -0,0 +1,202 @@
+# 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.
+
+"""NaFlex (NaViT + FlexiViT) preprocessing ops."""
+
+from big_vision.pp import utils
+from big_vision.pp.registry import Registry
+import big_vision.utils as u
+import tensorflow as tf
+
+
+def _get_image_size_for_seq(
+    image_hw,
+    patch_size: int,
+    max_sequence_len: int,
+    divisible_by_patch: bool = True,
+    eps: float = 1e-5):
+  """Determine scaling ratio and image size for `get_resize_to_sequence`.
+
+  Did not prove monotonicity necessary for binary search correctness, but this
+  works well in practice.
+
+  Args:
+    image_hw: Image height and width.
+    patch_size: Patchification patch size.
+    max_sequence_len: Maximum allowed sequence length for the resulting image.
+    divisible_by_patch: If True, the resulting image height and width must be
+      divisible by patch size.
+    eps: Small number used for binary search convergence.
+
+  Returns:
+    ratio: Scaling ratio to applied to image.
+    target_hw: Target image height and width taking into account the scaling
+      ratio and the `divisible_by_patch` constraint.
+  """
+  def search_not_done(lb, rb):
+    return (rb - lb) >= eps
+
+  def prepare_target_hw(ratio):
+    target_hw = tf.cast(image_hw, tf.float32) * ratio
+    if divisible_by_patch:
+      # Round to multiple of patch size as we want to avoid dropping patches.
+      target_hw = patch_size * tf.math.ceil(target_hw / patch_size)
+    # Ensure that the image is at least 1 patch in height / width.
+    target_hw = tf.maximum(target_hw, patch_size)
+    target_hw = tf.cast(target_hw, tf.int32)
+    return target_hw
+
+  def is_feasible(ratio):
+    target_hw = prepare_target_hw(ratio)
+    num_patches = target_hw / patch_size
+    sequence_len = tf.math.reduce_prod(num_patches)
+    return sequence_len <= max_sequence_len
+
+  def _search_fn(lb, rb):
+    mid = (lb + rb) / 2
+    return tf.cond(is_feasible(mid), lambda: (mid, rb), lambda: (lb, mid))
+
+  # Left and right boundaries for the binary search.
+  state = (tf.constant(eps / 10.), tf.constant(100.0))
+  ratio, _ = tf.while_loop(
+      search_not_done, _search_fn, state, parallel_iterations=1)
+  tf.assert_greater(
+      ratio, eps, message="Binary search failed - image too large?")
+  tf.assert_less(
+      ratio, 100.0, message="Binary search failed - image too small?")
+
+  return ratio, prepare_target_hw(ratio)
+
+
+@Registry.register("preprocess_ops.resize_to_sequence")
+@utils.InKeyOutKey(indefault="image", outdefault="image")
+def get_resize_to_sequence(
+    patch_size: int,
+    max_sequence_len: int,
+    divisible_by_patch: bool = True,
+    eps: float = 1e-5):
+  """Resizes image if it violates restrictions on sequence/side length.
+
+  This op attempts to resize the image in an AR-preserving manner such that:
+  - The sequence length of the resulting image (after patchification) is
+    maximized, but <= `max_sequence_len`.
+
+  This op *violates* the AR-preserving property if:
+  - Image size resulting from the above procedure is not a multiple of patch
+    size. In this case AR is distorted to ensure this condition is satisfied.
+
+  Args:
+    patch_size: Patchification patch size.
+    max_sequence_len: Maximum allowed sequence length for the resulting image.
+    divisible_by_patch: If True, the resulting image height and width must be
+      divisible by patch size.
+    eps: Small number used for binary search convergence.
+
+  Returns:
+    Pre-processing op.
+  """
+  def _resize_fn(image):
+    """Performs binary search to find a feasible image size."""
+    image_hw = tf.shape(image)[:2]
+    _, target_hw = _get_image_size_for_seq(
+        image_hw,
+        patch_size,
+        max_sequence_len,
+        divisible_by_patch=divisible_by_patch,
+        eps=eps)
+
+    # Actually resize image.
+    image = tf.image.resize(
+        image,
+        target_hw,
+        preserve_aspect_ratio=False,
+        antialias=True)
+    return tf.ensure_shape(image, [None, None, 3])
+  return _resize_fn
+
+
+@Registry.register("preprocess_ops.central_crop_to_sequence")
+@utils.InKeyOutKey(indefault="image", outdefault="image")
+def get_central_crop_to_sequence(
+    patch_size: int,
+    max_sequence_len: int,
+    divisible_by_patch: bool = True,
+    eps: float = 1e-5):
+  """Central crops image such that patch sequence length satisfies constraints.
+
+  Constraints used are the as in `resize_to_sequence`.
+
+  Args:
+    patch_size: Patchification patch size.
+    max_sequence_len: Maximum allowed sequence length for the resulting image.
+    divisible_by_patch: If True, the resulting image height and width must be
+      divisible by patch size.
+    eps: Small number used for binary search convergence.
+
+  Returns:
+    Pre-processing op.
+  """
+  def _central_crop_fn(image):
+    image_hw = tf.shape(image)[:2]
+    _, target_hw = _get_image_size_for_seq(
+        image_hw,
+        patch_size,
+        max_sequence_len,
+        divisible_by_patch=divisible_by_patch,
+        eps=eps)
+
+    tf.assert_greater(
+        image_hw + 1, target_hw,
+        "For central crop the image must be larger than target HW.")
+    offset_hw = (image_hw - target_hw) // 2
+    image = image[
+        offset_hw[0]:offset_hw[0] + target_hw[0],
+        offset_hw[1]:offset_hw[1] + target_hw[1],
+        :]
+    return tf.ensure_shape(image, [None, None, 3])
+  return _central_crop_fn
+
+
+@Registry.register("preprocess_ops.patchify")
+@utils.InKeyOutKey(indefault="image", outdefault="image")
+def get_patchify(patch_size):
+  """Reshapes image into patches and provides patch coordinates."""
+  ph, pw = utils.maybe_repeat(patch_size, 2)
+
+  def _patchify(img):
+    patches = tf.image.extract_patches(
+        img[None, ...], sizes=[1, ph, pw, 1], strides=[1, ph, pw, 1],
+        rates=[1, 1, 1, 1], padding="VALID")[0]
+    # Patches is now (nh, nw, ph*pw*3), i.e. contains flattened patches.
+    nh, nw, d = tf.shape(patches)[0], tf.shape(patches)[1], tf.shape(patches)[2]
+
+    # Get two (nh, nw) tensors of y/x indices of the patches.
+    gy, gx = tf.meshgrid(tf.range(nh), tf.range(nw), indexing="ij")
+
+    return {
+        "patches": tf.reshape(patches, (nh * nw, d)),
+        "yidx": tf.reshape(gy, [nh * nw]),
+        "xidx": tf.reshape(gx, [nh * nw]),
+        "type": tf.fill([nh * nw], 1),
+    }
+  return _patchify
+
+
+@Registry.register("preprocess_ops.tuplify")
+def get_tuplify(inkeys: list[str], outkey: str):
+  """Create a tuple of multiple inputs."""
+  def tuplify(data):
+    data[outkey] = tuple(u.tree_get(data, k) for k in inkeys)
+    return data
+  return tuplify

Tipsomaly/model/big_vision/tools/lit_demo/README.md

@@ -0,0 +1,26 @@
+# LiT-Demo
+
+See https://blog.tensorflow.org/2022/08/jax-on-web-with-tensorflowjs.html
+
+Demo originally appeared on Twitter
+https://twitter.com/AndreasPSteiner/status/1514722383818543106
+
+App published at
+https://google-research.github.io/vision_transformer/lit
+
+## Build
+
+Install packages (tested with node v16.17.0 and yarn 1.22.19)
+
+```bash
+yarn
+```
+
+
+## Run
+
+The web app will appear on http://localhost:8000
+
+```
+node build.js
+```