labelr 0.8.0__py3-none-any.whl → 0.10.0__py3-none-any.whl

labelr/apps/datasets.py

@@ -12,7 +12,11 @@ import typer
 from openfoodfacts import Flavor
 from openfoodfacts.utils import get_logger
 
-from labelr.export import export_from_ultralytics_to_hf
+from labelr.export.common import export_from_ultralytics_to_hf
+from labelr.export.object_detection import (
+    export_from_ls_to_hf_object_detection,
+    export_from_ls_to_ultralytics_object_detection,
+)
 
 from ..config import LABEL_STUDIO_DEFAULT_URL
 from ..types import ExportDestination, ExportSource, TaskType
@@ -99,7 +103,9 @@ def convert_object_detection_dataset(
     Studio format, and save it to a JSON file."""
     from datasets import load_dataset
 
-    from labelr.sample import format_object_detection_sample_from_hf_to_ls
+    from labelr.sample.object_detection import (
+        format_object_detection_sample_from_hf_to_ls,
+    )
 
     logger.info("Loading dataset: %s", repo_id)
     ds = load_dataset(repo_id)
@@ -207,10 +213,8 @@ def export(
     local files (ultralytics format)."""
     from label_studio_sdk.client import LabelStudio
 
-    from labelr.export import (
+    from labelr.export.object_detection import (
         export_from_hf_to_ultralytics_object_detection,
-        export_from_ls_to_hf_object_detection,
-        export_from_ls_to_ultralytics_object_detection,
     )
 
     if (to == ExportDestination.hf or from_ == ExportSource.hf) and repo_id is None:
@@ -303,3 +307,50 @@ def export(
                 is_openfoodfacts_dataset=is_openfoodfacts_dataset,
                 openfoodfacts_flavor=openfoodfacts_flavor,
             )
+
+
+@app.command()
+def export_llm_ds(
+    dataset_path: Annotated[
+        Path, typer.Option(..., help="Path to the JSONL dataset file")
+    ],
+    repo_id: Annotated[
+        str, typer.Option(..., help="Hugging Face Datasets repository ID to export to")
+    ],
+    split: Annotated[str, typer.Option(..., help="Dataset split to export")],
+    revision: Annotated[
+        str,
+        typer.Option(
+            help="Revision (branch, tag or commit) for the Hugging Face Datasets repository."
+        ),
+    ] = "main",
+    tmp_dir: Annotated[
+        Path | None,
+        typer.Option(
+            help="Path to a temporary directory to use for image processing",
+        ),
+    ] = None,
+    image_max_size: Annotated[
+        int | None,
+        typer.Option(
+            help="Maximum size (in pixels) for the images. If None, no resizing is performed.",
+        ),
+    ] = None,
+):
+    """Export LLM image extraction dataset with images only to Hugging Face
+    Datasets.
+    """
+    from labelr.export.llm import export_to_hf_llm_image_extraction
+    from labelr.sample.llm import load_llm_image_extraction_dataset_from_jsonl
+
+    sample_iter = load_llm_image_extraction_dataset_from_jsonl(
+        dataset_path=dataset_path
+    )
+    export_to_hf_llm_image_extraction(
+        sample_iter,
+        split=split,
+        repo_id=repo_id,
+        revision=revision,
+        tmp_dir=tmp_dir,
+        image_max_size=image_max_size,
+    )

labelr/apps/google_batch.py

@@ -0,0 +1,296 @@
+import asyncio
+import importlib
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+from google.genai.types import JSONSchema as GoogleJSONSchema
+from google.genai.types import Schema as GoogleSchema
+from openfoodfacts import Flavor
+from pydantic import BaseModel
+
+from labelr.google_genai import generate_batch_dataset, launch_batch_job
+
+app = typer.Typer()
+
+
+def convert_pydantic_model_to_google_schema(schema: type[BaseModel]) -> dict[str, Any]:
+    """Google doesn't support natively OpenAPI schemas, so we convert them to
+    Google `Schema` (a subset of OpenAPI)."""
+    return GoogleSchema.from_json_schema(
+        json_schema=GoogleJSONSchema.model_validate(schema.model_json_schema())
+    ).model_dump(mode="json", exclude_none=True, exclude_unset=True)
+
+
+@app.command()
+def generate_dataset(
+    data_path: Annotated[
+        Path,
+        typer.Option(
+            ...,
+            help="Path to a JSONL file containing the raw batch samples.",
+            exists=True,
+            dir_okay=False,
+            resolve_path=True,
+        ),
+    ],
+    output_path: Annotated[
+        Path,
+        typer.Option(
+            ...,
+            help="Path where to write the generated dataset file.",
+            exists=False,
+            dir_okay=False,
+            resolve_path=True,
+        ),
+    ],
+    config_module: Annotated[
+        str,
+        typer.Option(
+            ...,
+            help="Python module path (e.g., 'myschema') containing two variables: "
+            "OUTPUT_SCHEMA (a Pydantic class representing the output schema) and "
+            "INSTRUCTIONS (a str containing instructions to add before each sample).",
+        ),
+    ],
+    bucket_name: Annotated[
+        str,
+        typer.Option(
+            ...,
+            help="Name of the GCS bucket where the images are stored.",
+        ),
+    ] = "robotoff-batch",
+    bucket_dir_name: Annotated[
+        str,
+        typer.Option(
+            ...,
+            help="Directory name in the GCS bucket where the images are stored.",
+        ),
+    ] = "gemini-batch-images",
+    max_concurrent_uploads: Annotated[
+        int,
+        typer.Option(
+            ...,
+            help="Maximum number of concurrent uploads to GCS.",
+        ),
+    ] = 30,
+    base_image_dir: Annotated[
+        Path | None,
+        typer.Option(
+            ...,
+            help="Base directory to resolve local image paths from.",
+        ),
+    ] = None,
+    from_key: Annotated[
+        str | None,
+        typer.Option(
+            ...,
+            help="If specified, resume processing from this sample key.",
+        ),
+    ] = None,
+    skip_upload: Annotated[
+        bool, typer.Option(..., help="Skip uploading images to GCS")
+    ] = False,
+    thinking_level: Annotated[
+        str | None,
+        typer.Option(
+            ...,
+            help="Thinking level to use for the generation config.",
+        ),
+    ] = None,
+):
+    """Generate a dataset file in JSONL format to be used for batch
+    processing, using Gemini Batch Inference."""
+    typer.echo(f"Uploading images from '{data_path}' to GCS bucket '{bucket_name}'...")
+    typer.echo(f"Writing updated dataset to {output_path}...")
+    typer.echo(f"Max concurrent uploads: {max_concurrent_uploads}...")
+    typer.echo(f"Base image directory: {base_image_dir}...")
+    typer.echo(f"From key: {from_key}...")
+    typer.echo(f"Skip upload: {skip_upload}...")
+    typer.echo(f"Thinking level: {thinking_level}...")
+
+    module = importlib.import_module(config_module)
+    base_cls = getattr(module, "OUTPUT_SCHEMA")
+
+    if not issubclass(base_cls, BaseModel):
+        typer.echo(
+            f"Error: {config_module}.OUTPUT_SCHEMA is not a subclass of pydantic.BaseModel"
+        )
+        raise typer.Exit(code=1)
+
+    instructions = getattr(module, "INSTRUCTIONS", None) or None
+
+    if instructions:
+        typer.echo(f"Using instructions: '{instructions}'...")
+    else:
+        typer.echo("No instructions provided.")
+
+    # JSON Schema is supoorted natively by Vertex AI and Gemini APIs,
+    # but not yet on Batch Inference...
+    # So we convert the JSON schema to Google internal "Schema"
+    # google_json_schema = base_cls.model_json_schema()
+    google_json_schema = convert_pydantic_model_to_google_schema(base_cls)
+    asyncio.run(
+        generate_batch_dataset(
+            data_path=data_path,
+            output_path=output_path,
+            google_json_schema=google_json_schema,
+            instructions=instructions,
+            bucket_name=bucket_name,
+            bucket_dir_name=bucket_dir_name,
+            max_concurrent_uploads=max_concurrent_uploads,
+            base_image_dir=base_image_dir,
+            from_key=from_key,
+            skip_upload=skip_upload,
+            thinking_level=thinking_level,
+        )
+    )
+
+
+@app.command(name="launch-batch-job")
+def launch_batch_job_command(
+    run_name: Annotated[str, typer.Argument(..., help="Name of the batch job run")],
+    dataset_path: Annotated[Path, typer.Option(..., help="Path to the dataset file")],
+    model: Annotated[str, typer.Option(..., help="Model to use for the batch job")],
+    location: Annotated[
+        str,
+        typer.Option(..., help="GCP location where to run the batch job"),
+    ] = "europe-west4",
+):
+    """Launch a Gemini Batch Inference job."""
+    launch_batch_job(
+        run_name=run_name,
+        dataset_path=dataset_path,
+        model=model,
+        location=location,
+    )
+
+
+@app.command()
+def upload_training_dataset_from_predictions(
+    prediction_path: Annotated[
+        Path,
+        typer.Argument(
+            ...,
+            help="Path to the prediction JSONL file generated by Google Inference Batch",
+            exists=True,
+            dir_okay=False,
+            readable=True,
+        ),
+    ],
+    instructions_path: Annotated[
+        Path,
+        typer.Option(
+            ...,
+            help="Path to the file with the instruction prompt for the model",
+            exists=True,
+            dir_okay=False,
+            readable=True,
+        ),
+    ],
+    json_schema_path: Annotated[
+        Path,
+        typer.Option(
+            ...,
+            help="Path to the file with the JSON schema to follow",
+            dir_okay=False,
+            readable=True,
+        ),
+    ],
+    repo_id: Annotated[
+        str, typer.Option(help="Hugging Face Datasets repository ID to push to")
+    ],
+    revision: Annotated[
+        str,
+        typer.Option(
+            help="Revision (branch, tag or commit) to use for the Hugging Face Datasets repository"
+        ),
+    ] = "main",
+    is_openfoodfacts_dataset: Annotated[
+        bool, typer.Option(..., help="Whether this is an Open Food Facts dataset")
+    ] = False,
+    openfoodfacts_flavor: Annotated[
+        Flavor,
+        typer.Option(
+            ...,
+            help="Open Food Facts flavor of the dataset (if applicable)",
+        ),
+    ] = Flavor.off,
+    split: Annotated[str, typer.Option(..., help="Name of the split")] = "train",
+    tmp_dir: Annotated[
+        Path | None,
+        typer.Option(
+            ...,
+            help="Temporary directory to use for intermediate files, default to a temporary directory "
+            "generated automatically. This is useful to relaunch the command if it fails midway.",
+        ),
+    ] = None,
+    skip: Annotated[int, typer.Option(..., help="Number of samples to skip")] = 0,
+    limit: Annotated[
+        int | None,
+        typer.Option(
+            ..., help="Limit number of samples to process, or None for no limit"
+        ),
+    ] = None,
+    raise_on_invalid_sample: Annotated[
+        bool,
+        typer.Option(
+            ...,
+            help="Whether to raise an error on invalid samples instead of skipping them",
+        ),
+    ] = False,
+    image_max_size: Annotated[
+        int | None,
+        typer.Option(
+            help="Maximum size (in pixels) for the images. If None, no resizing is performed.",
+        ),
+    ] = None,
+):
+    """Upload a training dataset to a Hugging Face Datasets repository from a
+    Gemini batch prediction file."""
+    import tempfile
+
+    import orjson
+    from huggingface_hub import HfApi
+
+    from labelr.export.llm import export_to_hf_llm_image_extraction
+    from labelr.google_genai import generate_sample_iter
+
+    instructions = instructions_path.read_text()
+    print(f"Instructions: {instructions}")
+    json_schema = orjson.loads(json_schema_path.read_text())
+
+    api = HfApi()
+    config = {
+        "instructions": instructions,
+        "json_schema": json_schema,
+    }
+    with tempfile.TemporaryDirectory() as config_tmp_dir_str:
+        config_tmp_dir = Path(config_tmp_dir_str)
+        config_path = config_tmp_dir / "config.json"
+        config_path.write_text(
+            orjson.dumps(config, option=orjson.OPT_INDENT_2).decode("utf-8")
+        )
+        api.upload_file(
+            path_or_fileobj=config_path,
+            path_in_repo="config.json",
+            repo_id=repo_id,
+            repo_type="dataset",
+        )
+    sample_iter = generate_sample_iter(
+        prediction_path=prediction_path,
+        json_schema=json_schema,
+        is_openfoodfacts_dataset=is_openfoodfacts_dataset,
+        openfoodfacts_flavor=openfoodfacts_flavor,
+        skip=skip,
+        limit=limit,
+        raise_on_invalid_sample=raise_on_invalid_sample,
+    )
+    export_to_hf_llm_image_extraction(
+        sample_iter=sample_iter,
+        split=split,
+        repo_id=repo_id,
+        revision=revision,
+        tmp_dir=tmp_dir,
+        image_max_size=image_max_size,
+    )

labelr/apps/label_studio.py

@@ -398,7 +398,7 @@ def create_dataset_file(
     from openfoodfacts.images import extract_barcode_from_url, extract_source_from_url
     from openfoodfacts.utils import get_image_from_url
 
-    from labelr.sample import format_object_detection_sample_to_ls
+    from labelr.sample.object_detection import format_object_detection_sample_to_ls
 
     logger.info("Loading dataset: %s", input_file)
 

labelr/export/classification.py

@@ -0,0 +1,114 @@
+import functools
+import logging
+import pickle
+import tempfile
+from pathlib import Path
+
+import datasets
+from openfoodfacts.images import generate_image_url
+from openfoodfacts.types import Flavor
+from PIL import Image, ImageOps
+
+from labelr.export.common import _pickle_sample_generator
+from labelr.sample.classification import HF_DS_CLASSIFICATION_FEATURES
+
+logger = logging.getLogger(__name__)
+
+
+def export_from_ultralytics_to_hf_classification(
+    dataset_dir: Path,
+    repo_id: str,
+    label_names: list[str],
+    merge_labels: bool = False,
+    is_openfoodfacts_dataset: bool = False,
+    openfoodfacts_flavor: Flavor = Flavor.off,
+) -> None:
+    """Export an Ultralytics classification dataset to a Hugging Face dataset.
+
+    The Ultralytics dataset directory should contain 'train', 'val' and/or
+    'test' subdirectories, each containing subdirectories for each label.
+
+    Args:
+        dataset_dir (Path): Path to the Ultralytics dataset directory.
+        repo_id (str): Hugging Face repository ID to push the dataset to.
+        label_names (list[str]): List of label names.
+        merge_labels (bool): Whether to merge all labels into a single label
+            named 'object'.
+        is_openfoodfacts_dataset (bool): Whether the dataset is from
+            Open Food Facts. If True, the `off_image_id` and `image_url` will
+            be generated automatically. `off_image_id` is extracted from the
+            image filename.
+        openfoodfacts_flavor (Flavor): Flavor of Open Food Facts dataset. This
+            is ignored if `is_openfoodfacts_dataset` is False.
+    """
+    logger.info("Repo ID: %s, dataset_dir: %s", repo_id, dataset_dir)
+
+    if not any((dataset_dir / split).is_dir() for split in ["train", "val", "test"]):
+        raise ValueError(
+            f"Dataset directory {dataset_dir} does not contain 'train', 'val' or 'test' subdirectories"
+        )
+
+    # Save output as pickle
+    for split in ["train", "val", "test"]:
+        split_dir = dataset_dir / split
+
+        if not split_dir.is_dir():
+            logger.info("Skipping missing split directory: %s", split_dir)
+            continue
+
+        with tempfile.TemporaryDirectory() as tmp_dir_str:
+            tmp_dir = Path(tmp_dir_str)
+            for label_dir in (d for d in split_dir.iterdir() if d.is_dir()):
+                label_name = label_dir.name
+                if merge_labels:
+                    label_name = "object"
+                if label_name not in label_names:
+                    raise ValueError(
+                        "Label name %s not in provided label names (label names: %s)"
+                        % (label_name, label_names),
+                    )
+                label_id = label_names.index(label_name)
+
+                for image_path in label_dir.glob("*"):
+                    if is_openfoodfacts_dataset:
+                        image_stem_parts = image_path.stem.split("_")
+                        barcode = image_stem_parts[0]
+                        off_image_id = image_stem_parts[1]
+                        image_id = f"{barcode}_{off_image_id}"
+                        image_url = generate_image_url(
+                            barcode, off_image_id, flavor=openfoodfacts_flavor
+                        )
+                    else:
+                        image_id = image_path.stem
+                        barcode = ""
+                        off_image_id = ""
+                        image_url = ""
+                    image = Image.open(image_path)
+                    image.load()
+
+                    if image.mode != "RGB":
+                        image = image.convert("RGB")
+
+                    # Rotate image according to exif orientation using Pillow
+                    ImageOps.exif_transpose(image, in_place=True)
+                    sample = {
+                        "image_id": image_id,
+                        "image": image,
+                        "width": image.width,
+                        "height": image.height,
+                        "meta": {
+                            "barcode": barcode,
+                            "off_image_id": off_image_id,
+                            "image_url": image_url,
+                        },
+                        "category_id": label_id,
+                        "category_name": label_name,
+                    }
+                    with open(tmp_dir / f"{split}_{image_id}.pkl", "wb") as f:
+                        pickle.dump(sample, f)
+
+            hf_ds = datasets.Dataset.from_generator(
+                functools.partial(_pickle_sample_generator, tmp_dir),
+                features=HF_DS_CLASSIFICATION_FEATURES,
+            )
+            hf_ds.push_to_hub(repo_id, split=split)

labelr/export/common.py

@@ -0,0 +1,42 @@
+import pickle
+from pathlib import Path
+
+from openfoodfacts.types import Flavor
+
+from labelr.types import TaskType
+
+
+def _pickle_sample_generator(dir: Path):
+    """Generator that yields samples from pickles in a directory."""
+    for pkl in dir.glob("*.pkl"):
+        with open(pkl, "rb") as f:
+            yield pickle.load(f)
+
+
+def export_from_ultralytics_to_hf(
+    task_type: TaskType,
+    dataset_dir: Path,
+    repo_id: str,
+    label_names: list[str],
+    merge_labels: bool = False,
+    is_openfoodfacts_dataset: bool = False,
+    openfoodfacts_flavor: Flavor = Flavor.off,
+) -> None:
+    from labelr.export.classification import (
+        export_from_ultralytics_to_hf_classification,
+    )
+
+    if task_type != TaskType.classification:
+        raise NotImplementedError(
+            "Only classification task is currently supported for Ultralytics to HF export"
+        )
+
+    if task_type == TaskType.classification:
+        export_from_ultralytics_to_hf_classification(
+            dataset_dir=dataset_dir,
+            repo_id=repo_id,
+            label_names=label_names,
+            merge_labels=merge_labels,
+            is_openfoodfacts_dataset=is_openfoodfacts_dataset,
+            openfoodfacts_flavor=openfoodfacts_flavor,
+        )

labelr/export/llm.py

@@ -0,0 +1,91 @@
+import functools
+import logging
+import pickle
+import tempfile
+import typing
+from collections.abc import Iterator
+from pathlib import Path
+
+import datasets
+import tqdm
+from PIL import Image, ImageOps
+
+from labelr.export.common import _pickle_sample_generator
+from labelr.sample.llm import (
+    HF_DS_LLM_IMAGE_EXTRACTION_FEATURES,
+    LLMImageExtractionSample,
+)
+from labelr.utils import PathWithContext
+
+logger = logging.getLogger(__name__)
+
+
+def export_to_hf_llm_image_extraction(
+    sample_iter: Iterator[LLMImageExtractionSample],
+    split: str,
+    repo_id: str,
+    revision: str = "main",
+    tmp_dir: Path | None = None,
+    image_max_size: int | None = None,
+) -> None:
+    """Export LLM image extraction samples to a Hugging Face dataset.
+
+    Args:
+        sample_iter (Iterator[LLMImageExtractionSample]): Iterator of samples
+            to export.
+        split (str): Name of the dataset split (e.g., 'train', 'val').
+        repo_id (str): Hugging Face repository ID to push the dataset to.
+        revision (str): Revision (branch, tag or commit) to use for the
+            Hugging Face Datasets repository.
+        tmp_dir (Path | None): Temporary directory to use for intermediate
+            files. If None, a temporary directory will be created
+            automatically.
+        image_max_size (int | None): Maximum size (in pixels) for the images.
+    """
+    logger.info(
+        "Repo ID: %s, revision: %s, split: %s, tmp_dir: %s, image_max_size: %s",
+        repo_id,
+        revision,
+        split,
+        tmp_dir,
+        image_max_size,
+    )
+
+    tmp_dir_with_context: PathWithContext | tempfile.TemporaryDirectory
+    if tmp_dir:
+        tmp_dir.mkdir(parents=True, exist_ok=True)
+        tmp_dir_with_context = PathWithContext(tmp_dir)
+    else:
+        tmp_dir_with_context = tempfile.TemporaryDirectory()
+
+    with tmp_dir_with_context as tmp_dir_str:
+        tmp_dir = Path(tmp_dir_str)
+        for sample in tqdm.tqdm(sample_iter, desc="samples"):
+            image = sample.image
+            # Rotate image according to exif orientation using Pillow
+            image = typing.cast(Image.Image, ImageOps.exif_transpose(image))
+
+            if image_max_size is not None:
+                if image.height > image_max_size or image.width > image_max_size:
+                    image.thumbnail(
+                        (image_max_size, image_max_size),
+                        Image.Resampling.LANCZOS,
+                    )
+            image_id = sample.image_id
+            json_sample = {
+                "image_id": image_id,
+                "image": image,
+                "meta": {
+                    k: v for k, v in sample.meta.model_dump().items() if v is not None
+                },
+                "output": sample.output,
+            }
+            # Save output as pickle
+            with open(tmp_dir / f"{split}_{image_id}.pkl", "wb") as f:
+                pickle.dump(json_sample, f)
+
+        hf_ds = datasets.Dataset.from_generator(
+            functools.partial(_pickle_sample_generator, tmp_dir),
+            features=HF_DS_LLM_IMAGE_EXTRACTION_FEATURES,
+        )
+        hf_ds.push_to_hub(repo_id, split=split, revision=revision)