src/data/data_scripts/sa1b_cap.py

import json
from PIL import Image
import os
import re
from collections import defaultdict
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import urlparse

import datasets
from tqdm import tqdm
import dotenv
from ast import literal_eval
from git_utils.tsv_io import TSVFile
import subprocess
import tempfile
from contextlib import contextmanager
import hashlib
from datasets.utils.filelock import FileLock
import shutil


logger = datasets.logging.get_logger(__name__)

_CITATION = "TBD"

_DESCRIPTION = """\
SA1B, each mask region is annotated with a phrase describing the region.
the phrases are generated by GIT-2 model captioning masked objects on a
white background. 
"""
_HOMEPAGE = "TBD"
_LICENSE = "TBD"

_LATEST_VERSIONS = {
    "mask_region_descriptions": "0.0.1",
}

_BASE_IMAGE_METADATA_FEATURES = {
    "image_id": datasets.Value("int32"),
    "width": datasets.Value("int32"),
    "height": datasets.Value("int32"),
    "task_type": datasets.Value("string"),
}

_BASE_REGION_FEATURES = {
    "region_id": datasets.Value("int64"),
    "image_id": datasets.Value("int32"),
    "phrases": [datasets.Value("string")],
    "x": datasets.Value("int32"),
    "y": datasets.Value("int32"),
    "width": datasets.Value("int32"),
    "height": datasets.Value("int32"),
}


_BASE_MASK_FEATURES = {
    "size": [datasets.Value("int32")],
    "counts": datasets.Value("string"),
}

_BASE_MASK_REGION_FEATURES = {
    "region_id": datasets.Value("int64"),
    "image_id": datasets.Value("int32"),
    "phrases": [datasets.Value("string")],
    "x": datasets.Value("int32"),
    "y": datasets.Value("int32"),
    "width": datasets.Value("int32"),
    "height": datasets.Value("int32"),
    "mask": _BASE_MASK_FEATURES,
    # "area": datasets.Value("int32"),
    # "phrase_conf": datasets.Value("float32"),
}


_ANNOTATION_FEATURES = {
    "region_descriptions": {"regions": [_BASE_REGION_FEATURES]},
    "mask_region_descriptions": {"regions": [_BASE_MASK_REGION_FEATURES]},
}


class SA1BCapConfig(datasets.BuilderConfig):
    """BuilderConfig for SA1BCap."""

    def __init__(
        self,
        name: str,
        version: Optional[str] = None,
        with_image: bool = True,
        with_mask: bool = True,
        # 0
        sa1b_tar_url: Optional[str] = None,
        sa1b_tar_template: Optional[str] = None,
        # 1
        sa1b_annot_tsv_url: Optional[str] = None,
        sa1b_annot_template: Optional[str] = None,
        # 2
        sa1b_cap_tsv_url: Optional[str] = None,
        sa1b_cap_template: Optional[str] = None,
        # 3
        sa1b_filter_tsv_url: Optional[str] = None,
        sa1b_filter_template: Optional[str] = None,
        # 4
        sa1b_file_range: Optional[List[int]] = None,
        # 5
        training_args: Optional[Any] = None,
        # 6
        task_type: str = "caption",
        **kwargs,
    ):
        """BuilderConfig for SA1BCap.
        there should be **no dynamic** computation in __init__.
        The Config is first init in the DatasetBuilder constructor,
        then the attr here are to be modified in `load_dataset`.

        Args:
            name_version: name and version of the dataset.
            description: description of the dataset.
            image_dir: directory containing the images.
            annotation_dir: directory containing the annotations.
            **kwargs: keyword arguments forwarded to super.
        """
        _version = _LATEST_VERSIONS[name] if version is None else version
        # NOTE: f"{name}_v{_version}" is the param for `load_dataset`
        _name = f"{name}_v{_version}"
        super().__init__(version=datasets.Version(_version), name=_name, **kwargs)

        self._name_without_version = name

        # NOTE: the following attr can be overwritten by `load_dataset`
        self.with_image = with_image
        self.with_mask = with_mask

        self.sa1b_tar_url = sa1b_tar_url
        self.sa1b_tar_template = sa1b_tar_template

        self.sa1b_annot_tsv_url = sa1b_annot_tsv_url
        self.sa1b_annot_template = sa1b_annot_template

        self.sa1b_cap_tsv_url = sa1b_cap_tsv_url
        self.sa1b_cap_template = sa1b_cap_template

        self.sa1b_filter_tsv_url = sa1b_filter_tsv_url
        self.sa1b_filter_template = sa1b_filter_template

        self.sa1b_file_range = sa1b_file_range

        # NOTE: To determine whether it is main process or not
        self.training_args = training_args

        self.task_type = task_type

    @property
    def features(self):
        annoation_type = "mask_region_descriptions" if self.with_mask else "region_descriptions"
        logger.info(f"Using annotation type: {annoation_type} due to with_mask={self.with_mask}")
        return datasets.Features(
            {
                **({"image": datasets.Image()} if self.with_image else {}),
                **_BASE_IMAGE_METADATA_FEATURES,
                **_ANNOTATION_FEATURES[annoation_type],
            }
        )


class SA1BCap(datasets.GeneratorBasedBuilder):
    VERSION = datasets.Version("0.0.1")
    BUILDER_CONFIG_CLASS = SA1BCapConfig
    BUILDER_CONFIGS = [*[SA1BCapConfig(name="mask_region_descriptions", version=version) for version in ["0.0.1"]]]
    DEFAULT_CONFIG_NAME = "region_descriptions_v0.0.1"
    config: SA1BCapConfig

    def _info(self):
        return datasets.DatasetInfo(
            description=_DESCRIPTION,
            features=self.config.features,
            homepage=_HOMEPAGE,
            license=_LICENSE,
            citation=_CITATION,
            version=self.config.version,
        )

    def _split_generators(self, dl_manager: datasets.DownloadManager):
        sa1b_tar_url = self.config.sa1b_tar_url
        sa1b_annot_tsv_url = self.config.sa1b_annot_tsv_url
        sa1b_cap_tsv_url = self.config.sa1b_cap_tsv_url
        sa1b_filter_tsv_url = self.config.sa1b_filter_tsv_url

        sa1b_tar_template = self.config.sa1b_tar_template
        sa1b_annot_template = self.config.sa1b_annot_template
        sa1b_cap_template = self.config.sa1b_cap_template
        sa1b_filter_template = self.config.sa1b_filter_template

        sa1b_file_range = self.config.sa1b_file_range

        if sa1b_tar_url is None:
            raise ValueError("sa1b_tar_url is None")
        if sa1b_annot_tsv_url is None:
            raise ValueError("sa1b_annot_tsv_url is None")
        if sa1b_cap_tsv_url is None:
            raise ValueError("sa1b_cap_tsv_url is None")
        if sa1b_file_range is None:
            raise ValueError("sa1b_file_range is None. We need the exact file range to load the dataset.")

        try:
            sa1b_file_range = literal_eval(sa1b_file_range)
        except ValueError as e:
            sa1b_file_range = eval(sa1b_file_range)
        except Exception as e:
            logger.error(f"Failed to literal_eval sa1b_file_range: {e}")
            raise ValueError(f"Failed to literal_eval sa1b_file_range: {e}")

        _DL_URLS = {}

        # NOTE(xiaoke): load sas_key from .env
        logger.info(f"Try to load sas_key from .env file: {dotenv.load_dotenv('.env')}.")

        sa1b_tar_url_sas_key = os.getenv("SA1B_TAR_URL_SAS_KEY", None)
        if sa1b_tar_url_sas_key is None or os.path.exists(sa1b_tar_url):
            sa1b_tar_url_sas_key = ""
        _DL_URLS["sa1b_tar_urls"] = self._build_sa1b_urls(
            sa1b_tar_url, sa1b_tar_template, sa1b_file_range, sa1b_tar_url_sas_key
        )

        sa1b_annot_tsv_url_sas_key = os.getenv("SA1B_ANNOT_TSV_URL_SAS_KEY", None)
        if sa1b_annot_tsv_url_sas_key is None or os.path.exists(sa1b_annot_tsv_url):
            sa1b_annot_tsv_url_sas_key = ""
        _DL_URLS["sa1b_annot_tsv_urls"] = self._build_sa1b_urls(
            sa1b_annot_tsv_url, sa1b_annot_template, sa1b_file_range, sa1b_annot_tsv_url_sas_key
        )

        sa1b_cap_tsv_url_sas_key = os.getenv("SA1B_CAP_TSV_URL_SAS_KEY", None)
        if sa1b_cap_tsv_url_sas_key is None or os.path.exists(sa1b_cap_tsv_url):
            sa1b_cap_tsv_url_sas_key = ""
        _DL_URLS["sa1b_cap_tsv_urls"] = self._build_sa1b_urls(
            sa1b_cap_tsv_url, sa1b_cap_template, sa1b_file_range, sa1b_cap_tsv_url_sas_key
        )

        if sa1b_filter_tsv_url is None:
            logger.info(f"sa1b_filter_tsv_url is None, not filtering dataset.")
        else:
            sa1b_filter_tsv_url_sas_key = os.getenv("SA1B_FILTER_TSV_URL_SAS_KEY", None)
            if sa1b_filter_tsv_url_sas_key is None or os.path.exists(sa1b_filter_tsv_url):
                sa1b_filter_tsv_url_sas_key = ""
            _DL_URLS["sa1b_filter_tsv_urls"] = self._build_sa1b_urls(
                sa1b_filter_tsv_url, sa1b_filter_template, sa1b_file_range, sa1b_filter_tsv_url_sas_key
            )

        if dl_manager.is_streaming is False:
            raise ValueError("dl_manager.is_streaming is False. We need to stream the dataset. Because it is too big.")

        file_urls = dl_manager.download(_DL_URLS)
        num_tars = len(file_urls["sa1b_tar_urls"])
        self._num_tars = num_tars
        list_of_file_urls = []
        for num_tar in range(num_tars):
            list_of_file_urls.append(
                {
                    "sa1b_tar_url": file_urls["sa1b_tar_urls"][num_tar],
                    "sa1b_annot_tsv_url": file_urls["sa1b_annot_tsv_urls"][num_tar],
                    "sa1b_cap_tsv_url": file_urls["sa1b_cap_tsv_urls"][num_tar],
                    "sa1b_filter_tsv_url": file_urls["sa1b_filter_tsv_urls"][num_tar]
                    if "sa1b_filter_tsv_urls" in file_urls
                    else None,
                    "tar_idx": num_tar,
                }
            )

        return [
            datasets.SplitGenerator(
                name=datasets.Split.TRAIN,
                gen_kwargs={
                    "list_of_file_urls": list_of_file_urls,  # NOTE: It would be sharded as https://huggingface.co/docs/datasets/dataset_script#sharding, which would be much faster for downloading.
                    "iter_archive_func": dl_manager.iter_archive,
                },
            ),
        ]

    def _build_sa1b_urls(self, url, template, _range, sas_key):
        url_template = os.path.join(url, template)
        return [f"{url_template.format(i)}{sas_key}" for i in _range]

    def _generate_examples(self, list_of_file_urls, iter_archive_func):
        num_tars = len(list_of_file_urls)
        for i, one_file_urls in enumerate(list_of_file_urls):
            logger.info(f"Processing tar {one_file_urls['tar_idx']}/{self._num_tars}")
            tar_data_iter = self._process_one_tar(iter_archive_func, **one_file_urls)
            for image_id, data in tar_data_iter:
                yield image_id, data

    def _get_tsv_file(self, tsv_url):
        return TSVFile(tsv_url, open_func=open)

    def _process_one_tar(
        self,
        iter_archive_func,
        sa1b_tar_url,
        sa1b_annot_tsv_url,
        sa1b_cap_tsv_url,
        sa1b_filter_tsv_url=None,
        tar_idx=-1,
    ):
        # The `open` function of Python is extened with streaming loading from the Internet by `xopen` in `datasets.download.streaming_download_manager`.
        # After that, `xopen` is futher patched into `open` by `datasets.streaming`.

        sa1b_annot_tsv = self._get_tsv_file(sa1b_annot_tsv_url)

        sa1b_cap_tsv = self._get_tsv_file(sa1b_cap_tsv_url)

        sa1b_filter_tsv = None
        if sa1b_filter_tsv_url is not None:
            sa1b_filter_tsv = self._get_tsv_file(sa1b_filter_tsv_url)

        mapping_image_id_region_id_to_annot = self.build_mapping_image_id_region_id_to_annot(
            sa1b_annot_tsv, sa1b_cap_tsv, desc_prefix=f"[tar_idx={tar_idx}/{self._num_tars}]"
        )
        mapping_image_id_to_annots = self.build_mapping_image_id_to_annots(
            mapping_image_id_region_id_to_annot, desc_prefix=f"[tar_idx={tar_idx}/{self._num_tars}]"
        )
        del mapping_image_id_region_id_to_annot

        # NOTE: filter dataset if any:
        with TempFileForAzcopy(sa1b_tar_url) as _sa1b_tar_url:
            for name, buffer in iter_archive_func(_sa1b_tar_url):
                if name.endswith(".json"):
                    continue
                yield self._process_one_sample(name, buffer, mapping_image_id_to_annots)

    def _process_one_sample(self, name, buffer, mapping_image_id_to_annots):
        # name = './sa_%d.jpg"
        name = os.path.basename(name)
        image_id = int(name.split(".")[0].split("_")[-1])

        if self.config.with_image:
            # NOTE: check here see how hugging face datasets handle image
            # https://github.com/huggingface/datasets/blob/8b9649b3cfb49342e44873ce7e29e0c75eaf3efa/src/datasets/features/image.py#L130
            image = Image.open(buffer)
            image.load()
            if image.mode != "RGB":
                image = image.convert("RGB")
                logger.warning(f"convert {image_id} from {image.mode} to RGB")
            image_dict = dict(
                image=image,
                image_id=image_id,
                width=image.width,
                height=image.height,
            )
        else:
            image_dict = dict(
                image_id=image_id,
                width=-1,
                height=-1,
            )
        # convert to RGB is time consuming, from 5 it/s to 1it/s
        # image = image.convert("RGB")

        regions = mapping_image_id_to_annots[image_id]

        return image_id, dict(
            **image_dict,
            regions=regions,
            task_type=self.config.task_type,
        )

    def build_mapping_image_id_region_id_to_annot(self, annot_tsv, cap_tsv, desc_prefix=""):
        if len(annot_tsv) != len(cap_tsv):
            raise ValueError(
                f"len(annot_tsv) != len(cap_tsv): {len(annot_tsv)} != {len(cap_tsv)}. "
                f"Please check the data integrity for {annot_tsv} and {cap_tsv}."
            )

        # NOTE: Build index for fast retrieval of annoation.
        # This is compromised design as the tar file is extracted to image_id.json and image_id.jpg
        # NOTE: size: 965765982 bytes, 921.5 MB
        image_id_region_id_to_annot: Dict[int, Dict[int, List]] = defaultdict(dict)
        for cnt, (annot, cap) in enumerate(
            tqdm(
                zip(annot_tsv, cap_tsv),
                desc=f"{desc_prefix} building image_id_region_id_to_annot.",
                total=len(annot_tsv),
            )
        ):
            if annot[0] != cap[0]:
                raise ValueError(f"Cnt: {cnt}: annot[0] != cap[0], {annot[0]} != {cap[0]}, in {annot} != {cap}")

            # NOTE: identifier format is image_id-region_cnt-region_id
            image_id, region_cnt, region_id = list(map(int, cap[0].split("-")))

            annot_obj = json.loads(annot[1])  # Dict[str, Any], i.e. SA1B format
            # TODO: maybe update to other caption format
            cap_obj = json.loads(cap[1])  # NOTE: List[Dict[str, Any]], i.e. "caption" and "conf" from GIT2

            image_id_region_id_to_annot[image_id][region_id] = annot_obj
            image_id_region_id_to_annot[image_id][region_id]["captions"] = cap_obj

        return image_id_region_id_to_annot

    def build_mapping_image_id_to_annots(self, mapping_image_id_region_id_to_annot, desc_prefix):
        mapping_image_id_to_annots = {}
        for image_id, region_id_to_annot in tqdm(
            mapping_image_id_region_id_to_annot.items(),
            desc=f"{desc_prefix} building image_id_to_annots...",
            total=len(mapping_image_id_region_id_to_annot),
        ):
            annots = []
            for annot in region_id_to_annot.values():
                # _BASE_MASK_REGION_FEATURES
                region_id = annot["id"]
                image_id: int
                # TODO: maybe update to other caption format
                phrases = [caption["caption"] for caption in annot["captions"]]
                x, y, width, height = annot["bbox"]
                mask = annot["segmentation"]

                # Unused by model, but useful for filtering
                # phrase_conf = raw_annot["conf"]
                # area = raw_annot["area"]

                transformed_annot = dict(
                    region_id=region_id,
                    image_id=image_id,
                    phrases=phrases,
                    x=x,
                    y=y,
                    width=width,
                    height=height,
                    # area=area,
                    # phrase_conf=phrase_conf,
                )
                if self.config.with_mask:
                    transformed_annot["mask"] = mask
                annots.append(transformed_annot)

            mapping_image_id_to_annots[image_id] = annots
        return mapping_image_id_to_annots


class TempFileForAzcopy:
    def __init__(self, file_url):
        self.file_url = file_url
        self.temp_dir = self._get_temp_dir(file_url)
        self.temp_file = None
        self.lock_path = None

    def _get_lock_file_name(self, fname):
        path = urlparse(fname).path
        name = os.path.basename(path)
        return os.path.join(self.temp_dir, name), os.path.join(self.temp_dir, name + ".lock")

    def _get_temp_dir(self, fname):
        with tempfile.NamedTemporaryFile() as fp:
            base_temp_dir = os.path.dirname(fp.name)
        hash_str = hashlib.md5(fname.encode()).hexdigest()
        return os.path.join(base_temp_dir, "sa1b_cap-" + hash_str)

    def _is_file_open(self, file_path):
        return (
            subprocess.run(
                ["lsof", file_path],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            ).returncode
            == 0
        )

    def _remove_unopened_file(self, file_path):
        if self.temp_dir not in file_path:
            return

        logger.info("Try to remove file {}.".format(file_path))

        if self._is_file_open(file_path):
            logger.info(f"{file_path} is still open.")
        else:
            logger.info(f"{file_path} is all closed. So we remove it.")

            if os.path.exists(file_path):
                os.remove(file_path)
                logger.info(f"Successfully remove file {file_path}.")

            lock_file = file_path + ".lock"
            if os.path.exists(lock_file):
                os.remove(lock_file)
                logger.info(f"Successfully remove lock file {lock_file}.")

        if os.path.exists(self.temp_dir):
            if os.listdir(self.temp_dir) == 0:
                logger.info(f"{self.temp_dir} is not empty. So we do not remove it.")
            else:
                logger.info(f"Successfully remove temp dir {self.temp_dir} for {self.file_url}")
                shutil.rmtree(self.temp_dir, ignore_errors=True)

    def __enter__(self):
        has_azcopy = subprocess.run(["azcopy"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode
        has_azcopy = has_azcopy == 0
        file_url = self.file_url

        if "://" not in file_url:
            logger.debug("file_url is directory path.")
            return file_url
        if not has_azcopy:
            logger.warning("azcopy is not installed, skip using azcopy to prepare azure url.")
            return file_url

        if "blob.core.windows.net" not in file_url:
            logger.warning(f"file_url is not azure blob url, skip using azcopy to prepare azure url: {file_url}")
            return file_url

        temp_file, lock_path = self._get_lock_file_name(file_url)
        if not os.path.isdir(self.temp_dir):
            os.makedirs(self.temp_dir)
        with FileLock(lock_path):
            try:
                result = subprocess.run(
                    ["azcopy", "cp", file_url, temp_file],
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
                if result.returncode != 0:
                    raise ConnectionError(f"azcopy failed with return code {result.returncode}")
                logger.info(f"Successfully azcopy {file_url} to {temp_file}.")
                self.temp_file = temp_file
                self.lock_path = lock_path
                return temp_file

            except Exception as e:
                logger.error(f"azcopy failed with exception {e}. Use regular xopen instead which can be slow.")
                if os.path.isfile(temp_file):
                    os.remove(temp_file)
                if os.path.isfile(lock_path):
                    os.remove(lock_path)
                return file_url

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._remove_unopened_file(self.temp_file)

    def __del__(self):
        self._remove_unopened_file(self.temp_file)