Source code for segments.huggingface

from __future__ import annotations

import json
import logging
import os
import tempfile
from string import Template
from typing import TYPE_CHECKING, Any, Dict, cast

import requests
from PIL import Image
from segments.utils import load_image_from_url, load_label_bitmap_from_url

# https://adamj.eu/tech/2021/05/13/python-type-hints-how-to-fix-circular-imports/
if TYPE_CHECKING:
    from segments.typing import Release


#############
# Variables #
#############
logger = logging.getLogger(__name__)
try:
    import datasets
    from huggingface_hub import upload_file
except ImportError:
    logger.error(
        "Please install HuggingFace datasets first: pip install --upgrade datasets"
    )

# Add some functionality to the push_to_hub function of datasets.Dataset
push_to_hub_original = datasets.Dataset.push_to_hub


#############
# Functions #
#############
def push_to_hub(
    self: datasets.Dataset, repo_id: int, *args: Any, **kwargs: Any
) -> None:
    push_to_hub_original(self, repo_id, *args, **kwargs)

    # Upload the label file (https://huggingface.co/datasets/huggingface/label-files)
    if hasattr(self, "id2label"):
        # print("Uploading id2label.json")
        tmpfile = os.path.join(tempfile.gettempdir(), "id2label.json")
        with open(tmpfile, "w") as f:
            json.dump(self.id2label, f)

        upload_file(
            path_or_fileobj=tmpfile,
            path_in_repo="id2label.json",
            repo_id=f"datasets/{repo_id}",
        )

    # Upload README.md
    if hasattr(self, "readme"):
        # print("Uploading README.md")
        tmpfile = os.path.join(tempfile.gettempdir(), "README.md")
        with open(tmpfile, "w") as f:
            f.write(self.readme)

        upload_file(
            path_or_fileobj=tmpfile,
            path_in_repo="README.md",
            repo_id=f"datasets/{repo_id}",
        )


datasets.Dataset.push_to_hub = push_to_hub


def get_taxonomy_table(taxonomy: Dict[str, Any]) -> str:
    markdown_table = ""
    for category in taxonomy["categories"]:
        id_ = category["id"]
        name = category["name"]
        description = category["description"] if "description" in category else "-"
        markdown_table += f"| {id_} | {name} | {description} |\n"
    return markdown_table


[docs]def release2dataset(release: Release, download_images: bool = True) -> datasets.Dataset: """Create a Huggingface dataset from a release. Args: release: A Segments release resulting from :meth:`.get_release`. download_images: If images need to be downloaded from an AWS S3 url. Defaults to :obj:`True`. Returns: A HuggingFace dataset. Raises: :exc:`ValueError`: If the type of dataset is not yet supported. """ # try: # import datasets # except ImportError as e: # logger.error( # "Please install HuggingFace datasets first: pip install --upgrade datasets" # ) # raise e content = requests.get( cast(str, release.attributes.url) # TODO Fix in the backend. ) release_dict = json.loads(content.content) task_type = release_dict["dataset"]["task_type"] if task_type in ["vector", "bboxes", "keypoint"]: features = datasets.Features( { "name": datasets.Value("string"), "uuid": datasets.Value("string"), "image": {"url": datasets.Value("string")}, "status": datasets.Value("string"), "label": { "annotations": [ { "id": datasets.Value("int32"), "category_id": datasets.Value("int32"), "type": datasets.Value("string"), "points": [[datasets.Value("float32")]], } ], }, } ) elif task_type in ["segmentation-bitmap", "segmentation-bitmap-highres"]: features = datasets.Features( { "name": datasets.Value("string"), "uuid": datasets.Value("string"), "image": {"url": datasets.Value("string")}, "status": datasets.Value("string"), "label": { "annotations": [ { "id": datasets.Value("int32"), "category_id": datasets.Value("int32"), } ], "segmentation_bitmap": {"url": datasets.Value("string")}, }, } ) elif task_type in ["text-named-entities", "text-span-categorization"]: features = datasets.Features( { "name": datasets.Value("string"), "uuid": datasets.Value("string"), "text": datasets.Value("string"), "status": datasets.Value("string"), "label": { "annotations": [ { "start": datasets.Value("int32"), "end": datasets.Value("int32"), "category_id": datasets.Value("int32"), } ], }, } ) else: raise ValueError("This type of dataset is not yet supported.") samples = release_dict["dataset"]["samples"] data_rows = [] for sample in samples: try: del sample["labels"]["ground-truth"]["attributes"]["format_version"] except (KeyError, TypeError): pass data_row: Dict[str, Any] = {} # Name data_row["name"] = sample["name"] # Uuid data_row["uuid"] = sample["uuid"] # Status try: data_row["status"] = sample["labels"]["ground-truth"]["label_status"] except (KeyError, TypeError): data_row["status"] = "UNLABELED" # Image or text if task_type in [ "vector", "bboxes", "keypoint", "segmentation-bitmap", "segmentation-bitmap-highres", ]: try: data_row["image"] = sample["attributes"]["image"] except (KeyError, TypeError): data_row["image"] = {"url": None} elif task_type in ["text-named-entities", "text-span-categorization"]: try: data_row["text"] = sample["attributes"]["text"] except (KeyError, TypeError): data_row["text"] = None # Label try: label = sample["labels"]["ground-truth"]["attributes"] # Remove the image-level attributes if "attributes" in label: del label["attributes"] # Remove the object-level attributes for annotation in label["annotations"]: if "attributes" in annotation: del annotation["attributes"] data_row["label"] = label except (KeyError, TypeError): error_label: Dict[str, Any] = {"annotations": []} if task_type in ["segmentation-bitmap", "segmentation-bitmap-highres"]: error_label["segmentation_bitmap"] = {"url": None} data_row["label"] = error_label data_rows.append(data_row) # Now transform to column format dataset_dict: Dict[str, Any] = {key: [] for key in features.keys()} for data_row in data_rows: for key in dataset_dict.keys(): dataset_dict[key].append(data_row[key]) # Create the HF Dataset and flatten it dataset = datasets.Dataset.from_dict(dataset_dict, features, split="train") dataset = dataset.flatten() # Optionally download the images if ( task_type in [ "vector", "bboxes", "keypoint", "segmentation-bitmap", "segmentation-bitmap-highres", ] and download_images ): def download_image(data_row: Dict[str, Any]) -> Dict[str, Any]: try: data_row["image"] = load_image_from_url(data_row["image.url"]) except Exception: data_row["image"] = None return data_row def download_segmentation_bitmap(data_row: Dict[str, Any]) -> Dict[str, Any]: try: segmentation_bitmap = load_label_bitmap_from_url( data_row["label.segmentation_bitmap.url"] ) data_row["label.segmentation_bitmap"] = Image.fromarray( segmentation_bitmap ) except Exception: data_row["label.segmentation_bitmap"] = Image.new( "RGB", (1, 1) ) # TODO: replace with None return data_row dataset = dataset.map(download_image, remove_columns=["image.url"]) if task_type in ["segmentation-bitmap", "segmentation-bitmap-highres"]: dataset = dataset.map( download_segmentation_bitmap, remove_columns=["label.segmentation_bitmap.url"], ) # Reorder the features features = datasets.Features( { "name": dataset.features["name"], "uuid": dataset.features["uuid"], "status": dataset.features["status"], "image": datasets.Image(), "label.annotations": dataset.features["label.annotations"], "label.segmentation_bitmap": datasets.Image(), } ) dataset.info.features = features else: # Reorder the features features = datasets.Features( { "name": dataset.features["name"], "uuid": dataset.features["uuid"], "status": dataset.features["status"], "image": datasets.Image(), "label.annotations": dataset.features["label.annotations"], } ) dataset.info.features = features # Create id2label id2label = {} for category in release_dict["dataset"]["task_attributes"]["categories"]: id2label[category["id"]] = category["name"] id2label[0] = "unlabeled" dataset.id2label = id2label # Create readme.md and update DatasetInfo # https://stackoverflow.com/questions/6385686/is-there-a-native-templating-system-for-plain-text-files-in-python task_type = release_dict["dataset"]["task_type"] if task_type in ["segmentation-bitmap", "segmentation-bitmap-highres"]: task_category = "image-segmentation" elif task_type in ["vector", "bboxes"]: task_category = "object-detection" elif task_type in ["text-named-entities", "text-span-categorization"]: task_category = "named-entity-recognition" else: task_category = "other" info = { "name": release_dict["dataset"]["name"], "segments_url": f'https://segments.ai/{release_dict["dataset"]["owner"]}/{release_dict["dataset"]["name"]}', "short_description": release_dict["dataset"]["description"], "release": release_dict["name"], "taxonomy_table": get_taxonomy_table( release_dict["dataset"]["task_attributes"] ), "task_category": task_category, } # Create readme.md with open( os.path.join(os.path.dirname(__file__), "data", "dataset_card_template.md"), "r" ) as f: template = Template(f.read()) readme = template.substitute(info) dataset.readme = readme # Update DatasetInfo dataset.info.description = info["short_description"] dataset.info.homepage = info["segments_url"] return dataset