autolabel/utils.py

"""
utils.py — shared helpers for the autolabel pipeline.
"""

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif", ".webp"}


def collect_images(directory: Path) -> list[Path]:
    """Return sorted list of image files under *directory*."""
    images = sorted(
        p for p in directory.rglob("*") if p.suffix.lower() in IMAGE_EXTENSIONS
    )
    logger.info("Found %d image(s) in %s", len(images), directory)
    return images


def load_json(path: Path) -> Any:
    """Load and return JSON from *path*."""
    with path.open("r", encoding="utf-8") as fh:
        return json.load(fh)


def save_json(data: Any, path: Path, indent: int = 2) -> None:
    """Serialise *data* to JSON at *path*, creating parent dirs as needed."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=indent, ensure_ascii=False)
    logger.debug("Saved JSON → %s", path)


def detection_json_path(image_path: Path, detections_dir: Path) -> Path:
    """Return the expected detection JSON path for a given image."""
    return detections_dir / (image_path.stem + ".json")


def setup_logging(level: int = logging.INFO) -> None:
    """Configure root logger with a sensible format."""
    logging.basicConfig(
        level=level,
        format="%(asctime)s [%(levelname)s] %(name)s — %(message)s",
        datefmt="%H:%M:%S",
    )