src/edge_impulse.py

"""Upload a generated image dataset directly to an Edge Impulse project.

Uses the Edge Impulse ingestion REST API with a project API key. Image
filenames follow the ``label.<id>.jpg`` convention so Edge Impulse assigns
labels from the filename prefix automatically.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, List, Optional

import requests

_INGESTION_URL = "https://ingestion.edgeimpulse.com/api/{category}/files"
_TIMEOUT = 120
_BATCH_SIZE = 20
_IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".bmp")

ProgressFn = Callable[[str], None]


@dataclass
class UploadResult:
    uploaded: int = 0
    failed: int = 0
    errors: List[str] = field(default_factory=list)


def verify_api_key(api_key: str) -> bool:
    if not api_key or not api_key.strip():
        return False
    try:
        response = requests.post(
            _INGESTION_URL.format(category="training"),
            headers={"x-api-key": api_key.strip()},
            timeout=_TIMEOUT,
        )
    except requests.RequestException:
        return False
    return response.status_code not in (401, 403)


def _content_type(path: Path) -> str:
    return "image/png" if path.suffix.lower() == ".png" else "image/jpeg"


def _upload_batch(category: str, api_key: str, paths: List[Path], allow_duplicates: bool) -> tuple[int, Optional[str]]:
    headers = {"x-api-key": api_key}
    if not allow_duplicates:
        headers["x-disallow-duplicates"] = "1"

    files = []
    handles = []
    try:
        for path in paths:
            handle = path.open("rb")
            handles.append(handle)
            files.append(("data", (path.name, handle, _content_type(path))))
        response = requests.post(
            _INGESTION_URL.format(category=category),
            headers=headers,
            files=files,
            timeout=_TIMEOUT,
        )
    finally:
        for handle in handles:
            handle.close()

    if response.status_code == 200:
        return len(paths), None
    return 0, f"HTTP {response.status_code}: {response.text[:200]}"


def upload_dataset(
    dataset_dir: str,
    api_key: str,
    allow_duplicates: bool = False,
    progress: Optional[ProgressFn] = None,
) -> UploadResult:
    """Upload ``edge_impulse_upload/{training,testing}`` images to a project."""
    def log(message: str) -> None:
        if progress:
            progress(message)
        else:
            print(message)

    api_key = (api_key or "").strip()
    if not api_key:
        raise ValueError("An Edge Impulse API key is required to upload.")

    base = Path(dataset_dir) / "edge_impulse_upload"
    result = UploadResult()

    for split, category in (("training", "training"), ("testing", "testing")):
        split_dir = base / split
        if not split_dir.exists():
            continue
        imgs = sorted(p for p in split_dir.iterdir() if p.suffix.lower() in _IMAGE_EXTS)
        if not imgs:
            continue

        log(f"Uploading {len(imgs)} {category} image(s) to Edge Impulse...")
        for start in range(0, len(imgs), _BATCH_SIZE):
            batch = imgs[start:start + _BATCH_SIZE]
            uploaded, error = _upload_batch(category, api_key, batch, allow_duplicates)
            if error is None:
                result.uploaded += uploaded
                log(f"  {category}: {result.uploaded} uploaded")
            else:
                result.failed += len(batch)
                result.errors.append(f"{category} batch @ {start}: {error}")
                log(f"  WARNING: {category} batch failed: {error}")

    log(f"Edge Impulse upload complete: {result.uploaded} uploaded, {result.failed} failed.")
    return result