scripts/import_grocery_dataset.py

from __future__ import annotations

import argparse
import csv
import logging
import re
import shutil
import urllib.parse
import urllib.request
import zipfile
from pathlib import Path
from typing import Iterable


try:
    from _bootstrap import add_project_root_to_path
except ModuleNotFoundError:
    from scripts._bootstrap import add_project_root_to_path


ROOT = add_project_root_to_path()

from src.config import IMAGE_DIR


DEFAULT_REPO = "https://github.com/marcusklasson/GroceryStoreDataset"
DEFAULT_REF = "master"
DEFAULT_CACHE_DIR = ROOT / "data" / "_sources" / "grocery_store_dataset"
CATEGORY_LEVELS = {"super", "coarse", "fine"}
INVALID_CATEGORY_CHARS = re.compile(r'[<>:"/\\|?*]+')

logging.basicConfig(level=logging.INFO, format="%(levelname)s:%(name)s:%(message)s")
logger = logging.getLogger(__name__)


def _normalize_repo_url(repo_url: str) -> str:
    value = repo_url.strip().rstrip("/")
    if value.endswith(".git"):
        value = value[:-4]
    if "://" not in value:
        value = f"https://github.com/{value.lstrip('/')}"
    return value


def _parse_github_repo(repo_url: str) -> tuple[str, str, str]:
    normalized = _normalize_repo_url(repo_url)
    parsed = urllib.parse.urlparse(normalized)
    if not parsed.netloc:
        raise ValueError(f"Invalid repo URL: {repo_url}")
    parts = [part for part in parsed.path.split("/") if part]
    if len(parts) < 2:
        raise ValueError(f"Invalid repo URL: {repo_url}")
    owner, repo = parts[0], parts[1]
    base = f"{parsed.scheme or 'https'}://{parsed.netloc}/{owner}/{repo}"
    return owner, repo, base


def _build_zip_url(repo_base: str, ref: str) -> str:
    if ref.startswith("refs/"):
        return f"{repo_base}/archive/{ref}.zip"
    return f"{repo_base}/archive/refs/heads/{ref}.zip"


def _download_zip(url: str, dest: Path, force: bool) -> None:
    if dest.exists() and not force:
        logger.info("Using cached archive: %s", dest)
        return
    dest.parent.mkdir(parents=True, exist_ok=True)
    logger.info("Downloading: %s", url)
    request = urllib.request.Request(url, headers={"User-Agent": "image-retrieval-importer"})
    with urllib.request.urlopen(request) as response, open(dest, "wb") as handle:
        shutil.copyfileobj(response, handle)


def _extract_zip(zip_path: Path, dest_dir: Path, force: bool) -> Path:
    dest_dir.mkdir(parents=True, exist_ok=True)
    with zipfile.ZipFile(zip_path) as archive:
        top_levels = {
            Path(name).parts[0]
            for name in archive.namelist()
            if name and not name.endswith("/")
        }
    if len(top_levels) != 1:
        raise RuntimeError("Unexpected archive structure; could not locate repo root.")

    repo_root = dest_dir / next(iter(top_levels))
    if repo_root.exists() and not force:
        logger.info("Using extracted repo: %s", repo_root)
        return repo_root

    if repo_root.exists():
        shutil.rmtree(repo_root)

    logger.info("Extracting archive to: %s", dest_dir)
    with zipfile.ZipFile(zip_path) as archive:
        archive.extractall(dest_dir)
    return repo_root


def _iter_split_entries(split_file: Path) -> Iterable[str]:
    with split_file.open("r", encoding="utf-8") as handle:
        reader = csv.reader(handle)
        for row in reader:
            if not row:
                continue
            path_value = row[0].strip()
            if path_value:
                yield path_value


def _category_from_path(relative_path: str, level: str) -> str:
    parts = Path(relative_path).parts
    if level == "super":
        index = 1
    elif level == "coarse":
        index = 2
    else:
        index = 3
    if len(parts) > index:
        return parts[index]
    if len(parts) >= 2:
        return parts[-2]
    return "Uncategorized"


def _safe_category(name: str) -> str:
    cleaned = INVALID_CATEGORY_CHARS.sub("-", name.strip())
    return cleaned or "Uncategorized"


def _unique_destination(dest_dir: Path, filename: str, split: str) -> Path:
    candidate = dest_dir / filename
    if not candidate.exists():
        return candidate
    stem = Path(filename).stem
    suffix = Path(filename).suffix
    candidate = dest_dir / f"{stem}_{split}{suffix}"
    if not candidate.exists():
        return candidate
    counter = 2
    while True:
        candidate = dest_dir / f"{stem}_{split}_{counter}{suffix}"
        if not candidate.exists():
            return candidate
        counter += 1


def _resolve_path(value: Path) -> Path:
    return value if value.is_absolute() else ROOT / value


def _import_images(
    dataset_dir: Path,
    target_dir: Path,
    splits: list[str],
    category_level: str,
    dry_run: bool,
    limit: int | None,
) -> None:
    total_entries = 0
    copied = 0
    missing = 0
    skipped = 0
    seen: set[str] = set()

    for split in splits:
        split_file = dataset_dir / f"{split}.txt"
        if not split_file.exists():
            logger.warning("Split file not found: %s", split_file)
            continue

        for relative_path in _iter_split_entries(split_file):
            total_entries += 1
            relative_norm = relative_path.replace("\\", "/")
            if relative_norm in seen:
                skipped += 1
                continue
            seen.add(relative_norm)

            src = dataset_dir / relative_norm
            if not src.exists():
                missing += 1
                continue

            category = _safe_category(_category_from_path(relative_norm, category_level))
            dest_dir = target_dir / category
            if not dry_run:
                dest_dir.mkdir(parents=True, exist_ok=True)

            dest_path = _unique_destination(dest_dir, Path(relative_norm).name, split)
            if not dry_run:
                shutil.copy2(src, dest_path)
            copied += 1

            if limit and copied >= limit:
                logger.info("Reached limit of %d images", limit)
                logger.info("Processed entries: %d", total_entries)
                logger.info("Copied: %d, missing: %d, skipped: %d", copied, missing, skipped)
                return

    logger.info("Processed entries: %d", total_entries)
    logger.info("Copied: %d, missing: %d, skipped: %d", copied, missing, skipped)


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Download GroceryStoreDataset from GitHub and import images into data/images."
    )
    parser.add_argument("--repo", default=DEFAULT_REPO, help="GitHub repo URL or owner/repo.")
    parser.add_argument("--ref", default=DEFAULT_REF, help="Git ref (branch, tag, or commit).")
    parser.add_argument(
        "--dataset-dir",
        type=Path,
        help="Use an existing dataset folder (the one containing train.txt).",
    )
    parser.add_argument(
        "--cache-dir",
        type=Path,
        default=DEFAULT_CACHE_DIR,
        help="Where to download and extract the repo.",
    )
    parser.add_argument(
        "--target-dir",
        type=Path,
        default=IMAGE_DIR,
        help="Target images directory (default: data/images).",
    )
    parser.add_argument(
        "--splits",
        nargs="+",
        default=["train", "val", "test"],
        help="Which splits to import (train/val/test).",
    )
    parser.add_argument(
        "--category-level",
        choices=sorted(CATEGORY_LEVELS),
        default="coarse",
        help="Category granularity derived from the dataset paths.",
    )
    parser.add_argument("--limit", type=int, help="Import at most N images.")
    parser.add_argument("--dry-run", action="store_true", help="Scan without copying files.")
    parser.add_argument(
        "--force",
        action="store_true",
        help="Re-download and re-extract the dataset archive.",
    )
    parser.add_argument(
        "--cleanup",
        action="store_true",
        help="Remove the downloaded archive and extracted repo after import.",
    )
    args = parser.parse_args()

    splits = [split for split in args.splits if split in {"train", "val", "test"}]
    if not splits:
        raise SystemExit("No valid splits provided; use train, val, or test.")

    target_dir = _resolve_path(args.target_dir)

    if args.dataset_dir:
        dataset_dir = _resolve_path(args.dataset_dir)
        if dataset_dir.is_dir() and not (dataset_dir / "train.txt").exists():
            candidate = dataset_dir / "dataset"
            if candidate.exists():
                dataset_dir = candidate
        if not dataset_dir.exists():
            raise SystemExit(f"Dataset directory not found: {dataset_dir}")
        logger.info("Using existing dataset directory: %s", dataset_dir)
    else:
        cache_dir = _resolve_path(args.cache_dir)
        owner, repo, repo_base = _parse_github_repo(args.repo)
        zip_url = _build_zip_url(repo_base, args.ref)
        ref_slug = re.sub(r"[^A-Za-z0-9._-]+", "_", args.ref)
        zip_path = cache_dir / f"{repo}-{ref_slug}.zip"
        _download_zip(zip_url, zip_path, args.force)
        repo_root = _extract_zip(zip_path, cache_dir, args.force)
        dataset_dir = repo_root / "dataset"

    if not dataset_dir.exists():
        raise SystemExit(f"Dataset folder not found: {dataset_dir}")

    logger.info("Dataset directory: %s", dataset_dir)
    logger.info("Target directory: %s", target_dir)
    logger.info("Category level: %s", args.category_level)
    logger.info("Splits: %s", ", ".join(splits))

    _import_images(
        dataset_dir=dataset_dir,
        target_dir=target_dir,
        splits=splits,
        category_level=args.category_level,
        dry_run=args.dry_run,
        limit=args.limit,
    )

    if args.cleanup and not args.dataset_dir:
        cache_dir = _resolve_path(args.cache_dir)
        if cache_dir.exists():
            shutil.rmtree(cache_dir)
            logger.info("Removed cache directory: %s", cache_dir)


if __name__ == "__main__":
    main()