labelr 0.7.0__py3-none-any.whl → 0.9.0__py3-none-any.whl

labelr/evaluate/object_detection.py

@@ -0,0 +1,100 @@
+import tempfile
+from pathlib import Path
+
+import datasets
+import fiftyone as fo
+from huggingface_hub import hf_hub_download
+
+from labelr.dataset_features import OBJECT_DETECTION_DS_PREDICTION_FEATURES
+from labelr.utils import parse_hf_repo_id
+
+
+def convert_bbox_to_fo_format(
+    bbox: tuple[float, float, float, float],
+) -> tuple[float, float, float, float]:
+    # Bounding box coordinates should be relative values
+    # in [0, 1] in the following format:
+    # [top-left-x, top-left-y, width, height]
+    y_min, x_min, y_max, x_max = bbox
+    return (
+        x_min,
+        y_min,
+        (x_max - x_min),
+        (y_max - y_min),
+    )
+
+
+def visualize(
+    hf_repo_id: str,
+    dataset_name: str,
+    persistent: bool,
+):
+    hf_repo_id, hf_revision = parse_hf_repo_id(hf_repo_id)
+
+    file_path = hf_hub_download(
+        hf_repo_id,
+        filename="predictions.parquet",
+        revision=hf_revision,
+        repo_type="model",
+        # local_dir="./predictions/",
+    )
+    file_path = Path(file_path).absolute()
+    prediction_dataset = datasets.load_dataset(
+        "parquet",
+        data_files=str(file_path),
+        split="train",
+        features=OBJECT_DETECTION_DS_PREDICTION_FEATURES,
+    )
+    fo_dataset = fo.Dataset(name=dataset_name, persistent=persistent)
+
+    with tempfile.TemporaryDirectory() as tmpdir_str:
+        tmp_dir = Path(tmpdir_str)
+        for i, hf_sample in enumerate(prediction_dataset):
+            image = hf_sample["image"]
+            image_path = tmp_dir / f"{i}.jpg"
+            image.save(image_path)
+            split = hf_sample["split"]
+            sample = fo.Sample(
+                filepath=image_path,
+                split=split,
+                tags=[split],
+                image=hf_sample["image_id"],
+            )
+            ground_truth_detections = [
+                fo.Detection(
+                    label=hf_sample["objects"]["category_name"][i],
+                    bounding_box=convert_bbox_to_fo_format(
+                        bbox=hf_sample["objects"]["bbox"][i],
+                    ),
+                )
+                for i in range(len(hf_sample["objects"]["bbox"]))
+            ]
+            sample["ground_truth"] = fo.Detections(detections=ground_truth_detections)
+
+            if hf_sample["detected"] is not None and hf_sample["detected"]["bbox"]:
+                model_detections = [
+                    fo.Detection(
+                        label=hf_sample["detected"]["category_name"][i],
+                        bounding_box=convert_bbox_to_fo_format(
+                            bbox=hf_sample["detected"]["bbox"][i]
+                        ),
+                        confidence=hf_sample["detected"]["confidence"][i],
+                    )
+                    for i in range(len(hf_sample["detected"]["bbox"]))
+                ]
+                sample["model"] = fo.Detections(detections=model_detections)
+
+            fo_dataset.add_sample(sample)
+
+        # View summary info about the dataset
+        print(fo_dataset)
+
+        # Print the first few samples in the dataset
+        print(fo_dataset.head())
+
+        # Visualize the dataset in the FiftyOne App
+        session = fo.launch_app(fo_dataset)
+        fo_dataset.evaluate_detections(
+            "model", gt_field="ground_truth", eval_key="eval", compute_mAP=True
+        )
+        session.wait()

labelr/export.py

@@ -3,6 +3,7 @@ import logging
 import pickle
 import random
 import tempfile
+from collections.abc import Iterator
 from pathlib import Path
 
 import datasets
@@ -14,10 +15,13 @@ from PIL import Image, ImageOps
 
 from labelr.sample import (
     HF_DS_CLASSIFICATION_FEATURES,
+    HF_DS_LLM_IMAGE_EXTRACTION_FEATURES,
     HF_DS_OBJECT_DETECTION_FEATURES,
+    LLMImageExtractionSample,
     format_object_detection_sample_to_hf,
 )
 from labelr.types import TaskType
+from labelr.utils import PathWithContext
 
 logger = logging.getLogger(__name__)
 
@@ -77,13 +81,7 @@ def export_from_ls_to_hf_object_detection(
                 functools.partial(_pickle_sample_generator, tmp_dir),
                 features=HF_DS_OBJECT_DETECTION_FEATURES,
             )
-            hf_ds.push_to_hub(
-                repo_id,
-                split=split,
-                revision=revision,
-                # Create a PR if not pushing to main branch
-                create_pr=revision != "main",
-            )
+            hf_ds.push_to_hub(repo_id, split=split, revision=revision)
 
 
 def export_from_ls_to_ultralytics_object_detection(
@@ -461,3 +459,62 @@ def export_from_ultralytics_to_hf_classification(
                 features=HF_DS_CLASSIFICATION_FEATURES,
             )
             hf_ds.push_to_hub(repo_id, split=split)
+
+
+def export_to_hf_llm_image_extraction(
+    sample_iter: Iterator[LLMImageExtractionSample],
+    split: str,
+    repo_id: str,
+    revision: str = "main",
+    tmp_dir: Path | 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.
+    """
+    logger.info(
+        "Repo ID: %s, revision: %s, split: %s, tmp_dir: %s",
+        repo_id,
+        revision,
+        split,
+        tmp_dir,
+    )
+
+    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 = ImageOps.exif_transpose(image)
+            image_id = sample.image_id
+            sample = {
+                "image_id": image_id,
+                "image": image,
+                "meta": sample.meta.model_dump(),
+                "output": sample.output,
+            }
+            # Save output as pickle
+            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_LLM_IMAGE_EXTRACTION_FEATURES,
+        )
+        hf_ds.push_to_hub(repo_id, split=split, revision=revision)

labelr/google_genai.py

@@ -0,0 +1,415 @@
+import asyncio
+import mimetypes
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Literal
+from urllib.parse import urlparse
+
+import aiofiles
+import jsonschema
+import orjson
+import typer
+from gcloud.aio.storage import Storage
+from openfoodfacts import Flavor
+from openfoodfacts.images import download_image, generate_image_url
+from tqdm.asyncio import tqdm
+
+from labelr.sample import LLMImageExtractionSample, SampleMeta
+from labelr.utils import download_image_from_gcs
+
+try:
+    import google.genai  # noqa: F401
+except ImportError:
+    raise ImportError(
+        "The 'google-genai' package is required to use this module. "
+        "Please install labelr with the 'google' extra: "
+        "`pip install labelr[google]`"
+    )
+import aiohttp
+from google import genai
+from google.cloud import storage
+from google.genai.types import CreateBatchJobConfig, HttpOptions
+from google.genai.types import JSONSchema as GoogleJSONSchema
+from google.genai.types import Schema as GoogleSchema
+from openfoodfacts.types import JSONType
+from pydantic import BaseModel
+
+
+class RawBatchSamplePart(BaseModel):
+    type: Literal["text", "image"]
+    data: str
+
+
+class RawBatchSample(BaseModel):
+    key: str
+    parts: list[RawBatchSamplePart]
+    meta: JSONType = {}
+
+
+def convert_pydantic_model_to_google_schema(schema: type[BaseModel]) -> JSONType:
+    """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)
+
+
+async def download_image(url: str, session: aiohttp.ClientSession) -> bytes:
+    """Download an image from a URL and return its content as bytes.
+
+    Args:
+        url (str): URL of the image to download.
+    Returns:
+        bytes: Content of the downloaded image.
+    """
+    async with session.get(url) as response:
+        response.raise_for_status()
+        return await response.read()
+
+
+async def download_image_from_filesystem(url: str, base_dir: Path) -> bytes:
+    """Download an image from the filesystem and return its content as bytes.
+
+    Args:
+        url (str): URL of the image to download.
+        base_dir (Path): Base directory where images are stored.
+    Returns:
+        bytes: Content of the downloaded image.
+    """
+    file_path = urlparse(url).path[1:]  # Remove leading '/'
+    full_file_path = base_dir / file_path
+    async with aiofiles.open(full_file_path, "rb") as f:
+        return await f.read()
+
+
+async def upload_to_gcs(
+    image_url: str,
+    bucket_name: str,
+    blob_name: str,
+    session: aiohttp.ClientSession,
+    base_image_dir: Path | None = None,
+) -> dict:
+    """Upload data to Google Cloud Storage.
+    Args:
+        bucket_name (str): Name of the GCS bucket.
+        blob_name (str): Name of the blob (object) in the bucket.
+        data (bytes): Data to upload.
+        session (aiohttp.ClientSession): HTTP session to use for downloading
+            the image.
+        base_image_dir (Path | None): If provided, images will be read from
+            the filesystem under this base directory instead of downloading
+            them from their URLs.
+    Returns:
+        dict: Status of the upload operation.
+    """
+    if base_image_dir is None:
+        image_data = await download_image(image_url, session)
+    else:
+        image_data = await download_image_from_filesystem(image_url, base_image_dir)
+
+    client = Storage(session=session)
+
+    status = await client.upload(
+        bucket_name,
+        blob_name,
+        image_data,
+    )
+    return status
+
+
+async def upload_to_gcs_format_async(
+    sample: RawBatchSample,
+    google_json_schema: JSONType,
+    instructions: str | None,
+    bucket_name: str,
+    bucket_dir_name: str,
+    session: aiohttp.ClientSession,
+    base_image_dir: Path | None = None,
+    skip_upload: bool = False,
+    thinking_level: str | None = None,
+) -> JSONType | None:
+    parts: list[JSONType] = []
+
+    if instructions:
+        parts.append({"text": instructions})
+
+    for part in sample.parts:
+        if part.type == "image":
+            mime_type, _ = mimetypes.guess_type(part.data)
+            if mime_type is None:
+                raise ValueError(f"Cannot guess mimetype of file: {part.data}")
+
+            file_uri = part.data
+            image_blob_name = f"{bucket_dir_name}/{sample.key}/{Path(file_uri).name}"
+            # Download the image from the URL
+            if not skip_upload:
+                try:
+                    await upload_to_gcs(
+                        image_url=file_uri,
+                        bucket_name=bucket_name,
+                        blob_name=image_blob_name,
+                        session=session,
+                        base_image_dir=base_image_dir,
+                    )
+                except FileNotFoundError:
+                    return None
+
+            parts.append(
+                {
+                    "file_data": {
+                        "file_uri": f"gs://{bucket_name}/{image_blob_name}",
+                        "mime_type": mime_type,
+                    }
+                }
+            )
+        else:
+            parts.append({"text": part.data})
+
+    generation_config = {
+        "responseMimeType": "application/json",
+        "response_json_schema": google_json_schema,
+    }
+
+    if thinking_level is not None:
+        generation_config["thinkingConfig"] = {"thinkingLevel": thinking_level}
+
+    return {
+        "key": f"key:{sample.key}",
+        "request": {
+            "contents": [
+                {
+                    "parts": parts,
+                    "role": "user",
+                }
+            ],
+            "generationConfig": generation_config,
+        },
+    }
+
+
+async def generate_batch_dataset(
+    data_path: Path,
+    output_path: Path,
+    google_json_schema: JSONType,
+    instructions: str | None,
+    bucket_name: str,
+    bucket_dir_name: str,
+    max_concurrent_uploads: int = 30,
+    base_image_dir: Path | None = None,
+    from_key: str | None = None,
+    skip_upload: bool = False,
+    thinking_level: str | None = None,
+):
+    limiter = asyncio.Semaphore(max_concurrent_uploads)
+    ignore = True if from_key is None else False
+    missing_files = 0
+    async with aiohttp.ClientSession() as session:
+        async with asyncio.TaskGroup() as tg:
+            async with (
+                aiofiles.open(data_path, "r") as input_file,
+                aiofiles.open(output_path, "wb") as output_file,
+            ):
+                async with limiter:
+                    tasks = set()
+                    async for line in tqdm(input_file, desc="samples"):
+                        # print(f"line: {line}")
+                        sample = RawBatchSample.model_validate_json(line)
+                        # print(f"sample: {sample}")
+                        record_key = sample.key
+                        if from_key is not None and ignore:
+                            if record_key == from_key:
+                                ignore = False
+                            else:
+                                continue
+                        task = tg.create_task(
+                            upload_to_gcs_format_async(
+                                sample=sample,
+                                google_json_schema=google_json_schema,
+                                instructions=instructions,
+                                bucket_name=bucket_name,
+                                bucket_dir_name=bucket_dir_name,
+                                session=session,
+                                base_image_dir=base_image_dir,
+                                skip_upload=skip_upload,
+                                thinking_level=thinking_level,
+                            )
+                        )
+                        tasks.add(task)
+
+                        if len(tasks) >= max_concurrent_uploads:
+                            for task in tasks:
+                                await task
+                                updated_record = task.result()
+                                if updated_record is not None:
+                                    await output_file.write(
+                                        orjson.dumps(updated_record) + "\n".encode()
+                                    )
+                                else:
+                                    missing_files += 1
+                            tasks.clear()
+
+                    for task in tasks:
+                        await task
+                        updated_record = task.result()
+                        if updated_record is not None:
+                            await output_file.write(
+                                orjson.dumps(updated_record) + "\n".encode()
+                            )
+                        else:
+                            missing_files += 1
+
+    typer.echo(
+        f"Upload and dataset update completed. Wrote updated dataset to {output_path}. "
+        f"Missing files: {missing_files}."
+    )
+
+
+def launch_batch_job(
+    run_name: str,
+    dataset_path: Path,
+    model: str,
+    location: str,
+):
+    """Launch a Gemini Batch Inference job.
+
+    Args:
+        run_name (str): Name of the batch run.
+        dataset_path (Path): Path to the dataset file in JSONL format.
+        model (str): Model to use for the batch job. Example:
+            'gemini-2.5-flash'.
+        location (str): Location for the Vertex AI resources. Example:
+            'europe-west4'.
+    """
+    # We upload the dataset to a GCS bucket using the Gcloud
+
+    if model == "gemini-3-pro-preview" and location != "global":
+        typer.echo(
+            "Warning: only 'global' location is supported for 'gemini-3-pro-preview' model. Overriding location to 'global'."
+        )
+        location = "global"
+
+    storage_client = storage.Client()
+    bucket_name = "robotoff-batch"  # Replace with your bucket name
+    run_dir = f"gemini-batch/{run_name}"
+    input_file_blob_name = f"{run_dir}/inputs.jsonl"
+    bucket = storage_client.bucket(bucket_name)
+    blob = bucket.blob(input_file_blob_name)
+    blob.upload_from_filename(dataset_path)
+
+    client = genai.Client(
+        http_options=HttpOptions(api_version="v1"),
+        vertexai=True,
+        location=location,
+    )
+    output_uri = f"gs://{bucket_name}/{run_dir}"
+    job = client.batches.create(
+        model=model,
+        src=f"gs://{bucket_name}/{input_file_blob_name}",
+        config=CreateBatchJobConfig(dest=output_uri),
+    )
+    print(job)
+
+
+def generate_sample_iter(
+    prediction_path: Path,
+    json_schema: JSONType,
+    skip: int = 0,
+    limit: int | None = None,
+    is_openfoodfacts_dataset: bool = False,
+    openfoodfacts_flavor: Flavor = Flavor.off,
+    raise_on_invalid_sample: bool = False,
+) -> Iterator[LLMImageExtractionSample]:
+    """Generate training samples from a Gemini Batch Inference prediction
+    JSONL file.
+
+    Args:
+        prediction_path (Path): Path to the prediction JSONL file.
+        json_schema (JSONType): JSON schema to validate the predictions.
+        skip (int): Number of initial samples to skip.
+        limit (int | None): Maximum number of samples to generate.
+        is_openfoodfacts_dataset (bool): Whether the dataset is from Open Food
+            Facts.
+        openfoodfacts_flavor (Flavor): Flavor of the Open Food Facts dataset.
+    Yields:
+        Iterator[LLMImageExtractionSample]: Generated samples.
+    """
+    skipped = 0
+    invalid = 0
+    with prediction_path.open("r") as f_in:
+        for i, sample_str in enumerate(f_in):
+            if i < skip:
+                skipped += 1
+                continue
+            if limit is not None and i >= skip + limit:
+                break
+            sample = orjson.loads(sample_str)
+            try:
+                yield generate_sample_from_prediction(
+                    json_schema=json_schema,
+                    sample=sample,
+                    is_openfoodfacts_dataset=is_openfoodfacts_dataset,
+                    openfoodfacts_flavor=openfoodfacts_flavor,
+                )
+            except Exception as e:
+                if raise_on_invalid_sample:
+                    raise
+                else:
+                    typer.echo(
+                        f"Skipping invalid sample at line {i + 1} in {prediction_path}: {e}"
+                    )
+                    invalid += 1
+                    continue
+    if skipped > 0:
+        typer.echo(f"Skipped {skipped} samples.")
+    if invalid > 0:
+        typer.echo(f"Skipped {invalid} invalid samples.")
+
+
+def generate_sample_from_prediction(
+    json_schema: JSONType,
+    sample: JSONType,
+    is_openfoodfacts_dataset: bool = False,
+    openfoodfacts_flavor: Flavor = Flavor.off,
+) -> LLMImageExtractionSample:
+    """Generate a LLMImageExtractionSample from a prediction sample.
+    Args:
+        json_schema (JSONType): JSON schema to validate the predictions.
+        sample (JSONType): Prediction sample.
+        is_openfoodfacts_dataset (bool): Whether the dataset is from Open Food
+            Facts.
+        openfoodfacts_flavor (Flavor): Flavor of the Open Food Facts dataset.
+    Returns:
+        LLMImageExtractionSample: Generated sample.
+    """
+    image_id = sample["key"][len("key:") :]
+    response_str = sample["response"]["candidates"][0]["content"]["parts"][0]["text"]
+    image_uri = sample["request"]["contents"][0]["parts"][1]["file_data"]["file_uri"]
+    image = download_image_from_gcs(image_uri=image_uri)
+    response = orjson.loads(response_str)
+    jsonschema.validate(response, json_schema)
+
+    if is_openfoodfacts_dataset:
+        image_stem_parts = image_id.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_id
+        barcode = ""
+        off_image_id = ""
+        image_url = ""
+
+    sample_meta = SampleMeta(
+        barcode=barcode,
+        off_image_id=off_image_id,
+        image_url=image_url,
+    )
+    return LLMImageExtractionSample(
+        image_id=image_id,
+        image=image,
+        output=orjson.dumps(response).decode("utf-8"),
+        meta=sample_meta,
+    )

labelr/main.py

@@ -4,9 +4,11 @@ import typer
 from openfoodfacts.utils import get_logger
 
 from labelr.apps import datasets as dataset_app
-from labelr.apps import projects as project_app
+from labelr.apps import evaluate as evaluate_app
+from labelr.apps import google_batch as google_batch_app
+from labelr.apps import hugging_face as hf_app
+from labelr.apps import label_studio as ls_app
 from labelr.apps import train as train_app
-from labelr.apps import users as user_app
 
 app = typer.Typer(pretty_exceptions_show_locals=False)
 
@@ -58,22 +60,35 @@ def predict(
         typer.echo(result)
 
 
-app.add_typer(user_app.app, name="users", help="Manage Label Studio users")
 app.add_typer(
-    project_app.app,
-    name="projects",
-    help="Manage Label Studio projects (create, import data, etc.)",
+    ls_app.app,
+    name="ls",
+    help="Manage Label Studio projects (create, import data, etc.).",
+)
+app.add_typer(
+    hf_app.app,
+    name="hf",
+    help="Manage Hugging Face Datasets repositories.",
 )
 app.add_typer(
     dataset_app.app,
     name="datasets",
     help="Manage datasets (convert, export, check, etc.)",
 )
-
 app.add_typer(
     train_app.app,
     name="train",
-    help="Train models",
+    help="Train models.",
+)
+app.add_typer(
+    evaluate_app.app,
+    name="evaluate",
+    help="Visualize and evaluate trained models.",
+)
+app.add_typer(
+    google_batch_app.app,
+    name="google-batch",
+    help="Generate datasets and launch batch jobs on Google Gemini.",
 )
 
 if __name__ == "__main__":