ivelin/ui_refexp

1# coding=utf-82# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8#     http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""Dataset script for UI Referring Expressions based on the UIBert RefExp dataset."""16 17 18import csv19import glob20import os21import tensorflow as tf22import re23import datasets24import json25import numpy as np26 27# Find for instance the citation on arxiv or on the dataset repo/website28_CITATION = """\29@misc{bai2021uibert,30      title={UIBert: Learning Generic Multimodal Representations for UI Understanding},31      author={Chongyang Bai and Xiaoxue Zang and Ying Xu and Srinivas Sunkara and Abhinav Rastogi and Jindong Chen and Blaise Aguera y Arcas},32      year={2021},33      eprint={2107.13731},34      archivePrefix={arXiv},35      primaryClass={cs.CV}36}37"""38 39# TODO: Add description of the dataset here40# You can copy an official description41_DESCRIPTION = """\42This dataset is intended for UI understanding, referring expression and action automation model training. It's based on the UIBert RefExp dataset from Google Research, which is based on the RICO dataset.43"""44 45# TODO: Add a link to an official homepage for the dataset here46_HOMEPAGE = "https://github.com/google-research-datasets/uibert"47 48# TODO: Add the licence for the dataset here if you can find it49_LICENSE = "CC BY 4.0"50 51# Add link to the official dataset URLs here52# The HuggingFace dataset library don't host the datasets but only point to the original files53# This can be an arbitrary nested dict/list of URLs (see below in `_split_generators` method)54_DATA_URLs = {55    "ui_refexp": "https://storage.googleapis.com/crowdstf-rico-uiuc-4540/rico_dataset_v0.1/unique_uis.tar.gz"56    # "https://huggingface.co/datasets/ncoop57/rico_captions/resolve/main/captions_hierarchies_images_filtered.zip",57}58 59_METADATA_URLS = {60    "ui_refexp": {61        "train": "https://github.com/google-research-datasets/uibert/raw/main/ref_exp/train.tfrecord",62        "validation": "https://github.com/google-research-datasets/uibert/raw/main/ref_exp/dev.tfrecord",63        "test": "https://github.com/google-research-datasets/uibert/raw/main/ref_exp/test.tfrecord"64    }65}66 67 68def tfrecord2list(tfr_file: None):69    """Filter and convert refexp tfrecord file to a list of dict object.70    Each sample in the list is a dict with the following keys: (image_id, prompt, target_bounding_box)"""71    raw_tfr_dataset = tf.data.TFRecordDataset([tfr_file])72    count = 073    donut_refexp_dict = []74    for raw_record in raw_tfr_dataset:75        count += 176        example = tf.train.Example()77        example.ParseFromString(raw_record.numpy())78        # print(f"total UI objects in this sample: {len(example.features.feature['image/object/bbox/xmin'].float_list.value)}")79        # print(f"feature keys: {example.features.feature.keys}")80        donut_refexp = {}81        image_id = example.features.feature['image/id'].bytes_list.value[0].decode()82        donut_refexp["image_id"] = image_id83        donut_refexp["prompt"] = example.features.feature["image/ref_exp/text"].bytes_list.value[0].decode()84        object_idx = example.features.feature["image/ref_exp/label"].int64_list.value[0]85        object_idx = int(object_idx)86        # print(f"object_idx: {object_idx}")87        object_bb = {}88        # print(f"example.features.feature['image/object/bbox/xmin']: {example.features.feature['image/object/bbox/xmin'].float_list.value[object_idx]}")89        object_bb["xmin"] = example.features.feature['image/object/bbox/xmin'].float_list.value[object_idx]90        object_bb["ymin"] = example.features.feature['image/object/bbox/ymin'].float_list.value[object_idx]91        object_bb["xmax"] = example.features.feature['image/object/bbox/xmax'].float_list.value[object_idx]92        object_bb["ymax"] = example.features.feature['image/object/bbox/ymax'].float_list.value[object_idx]93        donut_refexp["target_bounding_box"] = object_bb94        donut_refexp_dict.append(donut_refexp)95        if count != 3:96            continue97        print(f"Donut refexp: {donut_refexp}")98        # for key, feature in example.features.feature.items():99        #   if key in ['image/id', "image/ref_exp/text", "image/ref_exp/label", 'image/object/bbox/xmin', 'image/object/bbox/ymin', 'image/object/bbox/xmax', 'image/object/bbox/ymax']:100        #     print(key, feature)101 102    print(f"Total samples in the raw dataset: {count}")103    return donut_refexp_dict104 105 106class UIRefExp(datasets.GeneratorBasedBuilder):107    """Dataset with (image, question, answer) fields derive from UIBert RefExp."""108 109    VERSION = datasets.Version("1.1.0")110 111    # This is an example of a dataset with multiple configurations.112    # If you don't want/need to define several sub-sets in your dataset,113    # just remove the BUILDER_CONFIG_CLASS and the BUILDER_CONFIGS attributes.114 115    # If you need to make complex sub-parts in the datasets with configurable options116    # You can create your own builder configuration class to store attribute, inheriting from datasets.BuilderConfig117    # BUILDER_CONFIG_CLASS = MyBuilderConfig118 119    # You will be able to load one or the other configurations in the following list with120    # data = datasets.load_dataset('my_dataset', 'first_domain')121    # data = datasets.load_dataset('my_dataset', 'second_domain')122    BUILDER_CONFIGS = [123        datasets.BuilderConfig(124            name="ui_refexp",125            version=VERSION,126            description="Contains 66k+ unique UI screens. For each UI, we present a screenshot (JPG file) and the text shown on the screen that was extracted using an OCR model.",127        )128        #    ,129        #     # datasets.BuilderConfig(130        #     #     name="screenshots_captions_filtered",131        #     #     version=VERSION,132        #     #     description="Contains 25k unique UI screens. For each UI, we present a screenshot (JPG file) and the text shown on the screen that was extracted using an OCR model. Filtering was done as discussed in this paper: https://aclanthology.org/2020.acl-main.729.pdf",133        #     # ),134    ]135 136    DEFAULT_CONFIG_NAME = "ui_refexp"137 138    def _info(self):139        features = datasets.Features(140            {141                "image": datasets.Image(),142                "image_id": datasets.Value("string"),143                "image_file_path": datasets.Value("string"),144                # click the search button next to menu drawer at the top of the screen145                "prompt": datasets.Value("string"),146                # json: {xmin, ymin, xmax, ymax}, normalized screen reference values between 0 and 1147                "target_bounding_box": datasets.Value("string"),148            }149        )150 151        return datasets.DatasetInfo(152            description=_DESCRIPTION,153            features=features,154            homepage=_HOMEPAGE,155            license=_LICENSE,156            citation=_CITATION,157        )158 159    def _split_generators(self, dl_manager):160        """Returns SplitGenerators."""161        # This method is tasked with downloading/extracting the data and defining the splits depending on the configuration162        # If several configurations are possible (listed in BUILDER_CONFIGS), the configuration selected by the user is in self.config.name163 164        # dl_manager is a datasets.download.DownloadManager that can be used to download and extract URLs165        # It can accept any type or nested list/dict and will give back the same structure with the url replaced with path to local files.166        # By default the archives will be extracted and a path to a cached folder where they are extracted is returned instead of the archive167        # download and extract TFRecord labeling metadata168        local_tfrs = {}169        for split, tfrecord_url in _METADATA_URLS[self.config.name].items():170            local_tfr_file = dl_manager.download(tfrecord_url)171            local_tfrs[split] = local_tfr_file172        # download image files173        image_urls = _DATA_URLs[self.config.name]174        archive_path = dl_manager.download(image_urls)175 176        return [177            datasets.SplitGenerator(178                name=datasets.Split.TRAIN,179                # These kwargs will be passed to _generate_examples180                gen_kwargs={181                    "metadata_file": local_tfrs["train"],182                    "images": dl_manager.iter_archive(archive_path),183                    "split": "train",184 185                },186            ),187            datasets.SplitGenerator(188                name=datasets.Split.VALIDATION,189                # These kwargs will be passed to _generate_examples190                gen_kwargs={191                    "metadata_file": local_tfrs["validation"],192                    "images": dl_manager.iter_archive(archive_path),193                    "split": "validation",194                },195            ),196            datasets.SplitGenerator(197                name=datasets.Split.TEST,198                # These kwargs will be passed to _generate_examples199                gen_kwargs={200                    "metadata_file": local_tfrs["test"],201                    "images": dl_manager.iter_archive(archive_path),202                    "split": "test",203                },204            )205        ]206 207    def _generate_examples(208        self,209        metadata_file,210        images,211        split,  # method parameters are unpacked from `gen_kwargs` as given in `_split_generators`212    ):213        """Yields examples as (key, example) tuples."""214        # This method handles input defined in _split_generators to yield (key, example) tuples from the dataset.215        # The `key` is here for legacy reason (tfds) and is not important in itself.216        # filter tfrecord and convert to json217 218        metadata = tfrecord2list(metadata_file)219        files_to_keep = set()220        image_labels = {}221        for sample in metadata:222            image_id = sample["image_id"]223            files_to_keep.add(image_id)224            labels = image_labels.get(image_id)225            if isinstance(labels, list):226                labels.append(sample)227            else:228                labels = [sample]229                image_labels[image_id] = labels230        _id = 0231        for file_path, file_obj in images:232            image_id = re.search("(\d+).jpg", file_path)233            if image_id:234                image_id = image_id.group(1)235                if image_id in files_to_keep:236                    image_bytes = file_obj.read()237                    for labels in image_labels[image_id]:238                        bb_json = json.dumps(labels["target_bounding_box"])239                        yield _id, {240                            "image": {"path": file_path, "bytes": image_bytes},241                            "image_id": image_id,242                            "image_file_path": file_path,243                            "prompt": labels["prompt"],244                            "target_bounding_box": bb_json245                        }246                        _id += 1247