ml/scripts/url_cache_prefetch.py

# -*- coding: utf-8 -*-
"""
Prefetch downloader for TACO manifests:
- Reads "<file_name>\t<url>" (or comma CSV) from --manifest
- Downloads images into a URL-based cache (--cache_dir; hashed filename with image suffix)
- Links/copies them into --dataset_images_dir (per split)
Improvements:
- Parallel downloads (--workers)
- Optional HEAD checks (Content-Type/Status) before GET (--no-head-check to disable)
- Image sanity (PIL.verify) to filter broken responses (--no-image-verify to disable)
- Exponential backoff in http_get (retries)
- Missing report as CSV next to the manifest
- Link mode: symlink (Unix) or copy (Windows default) via --link-mode {auto,copy,symlink}
- Limit/shuffle preserved
"""

import os
import sys
import io
import csv
import time
import hashlib
import argparse
import random
from pathlib import Path
from typing import List, Tuple, Optional
from urllib.parse import urlparse
from concurrent.futures import ThreadPoolExecutor, as_completed
from collections import defaultdict

import requests
from PIL import Image


# --------------------------- Path/manifest helpers ---------------------------

def url_to_cache_path(url: str, cache_root: Path) -> Path:
    h = hashlib.sha1(url.encode("utf-8")).hexdigest()[:20]
    # Derive suffix from URL, else default to .jpg
    suf = ""
    for s in [".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"]:
        if url.lower().split("?")[0].endswith(s):
            suf = s
            break
    if not suf:
        suf = ".jpg"
    return cache_root / f"{h}{suf}"


def load_manifest(path: Path) -> List[Tuple[str, str]]:
    """
    Expected lines: "<file_name>\t<url>".
    Falls back to comma CSV if no tab is present.
    Ignores empty and malformed lines.
    """
    out: List[Tuple[str, str]] = []
    txt = path.read_text(encoding="utf-8")
    # Erlaubt gemischte Delimiter (Tab bevorzugt)
    for raw in txt.splitlines():
        line = raw.strip()
        if not line:
            continue
        parts: List[str]
        if "\t" in line:
            parts = [p.strip() for p in line.split("\t")]
        else:
            # CSV-Fallback (einfach): nur 2 Spalten relevant
            parts = [p.strip() for p in line.split(",")]
        if len(parts) < 2:
            continue
        file_name, url = parts[0], parts[1]
        if not file_name or not url:
            continue
        out.append((file_name, url))
    return out


# --------------------------- Network/download helpers ------------------------

def domain_of(url: str) -> str:
    try:
        return urlparse(url).netloc.lower()
    except Exception:
        return "unknown"


def head_ok(url: str, timeout: int = 10) -> bool:
    """
    Quick HEAD validation:
    - HTTP status < 400
    - Content-Type contains 'image' (or empty = accept)
    """
    try:
        r = requests.head(url, timeout=timeout, allow_redirects=True)
        if r.status_code >= 400:
            return False
        ct = r.headers.get("Content-Type", "")
        return ("image" in ct.lower()) or (ct.strip() == "")
    except Exception:
        return False


def http_get_verified(
    url: str,
    dst: Path,
    retries: int = 3,
    timeout: int = 25,
    verify_image: bool = True,
    backoff_base: float = 0.6,
) -> Tuple[bool, str]:
    """
    GET with retries + optional image verification (PIL.verify()).
    Writes to 'dst' only on success. Returns (ok, errmsg).
    """
    last_err = ""
    for attempt in range(retries):
        try:
            r = requests.get(url, timeout=timeout, allow_redirects=True)
            r.raise_for_status()
            data = r.content

            if verify_image:
                try:
                    Image.open(io.BytesIO(data)).verify()
                except Exception as e:
                    last_err = f"invalid image bytes: {e}"
                    # Backoff vor erneutem Versuch
                    time.sleep(backoff_base * (attempt + 1))
                    continue

            dst.parent.mkdir(parents=True, exist_ok=True)
            with dst.open("wb") as f:
                f.write(data)
            return True, ""
        except Exception as e:
            last_err = str(e)
            time.sleep(backoff_base * (attempt + 1))
    return False, last_err or "unknown error"


# ------------------------------- Linking ------------------------------------

def link_into_dataset(cache_path: Path, dataset_img_path: Path, link_mode: str = "auto") -> None:
    """
    Link/copy from cache into the dataset path.
    link_mode:
        - auto    -> Windows = copy, otherwise symlink
        - copy    -> always copy
        - symlink -> always symlink (may fail on Windows without privileges)
    Existing destination file is replaced (Unix: unlink + symlink).
    """
    dataset_img_path.parent.mkdir(parents=True, exist_ok=True)

    effective_mode = link_mode
    if link_mode == "auto":
        effective_mode = "copy" if os.name == "nt" else "symlink"

    if effective_mode == "copy":
        if not dataset_img_path.exists():
            import shutil
            shutil.copy2(cache_path, dataset_img_path)
        else:
            # destination exists -> overwrite
            import shutil
            shutil.copy2(cache_path, dataset_img_path)
    else:
        # symlink
        if dataset_img_path.exists() or dataset_img_path.is_symlink():
            try:
                dataset_img_path.unlink()
            except FileNotFoundError:
                pass
        dataset_img_path.symlink_to(cache_path.resolve())


# ------------------------------- Worker -------------------------------------

def process_one(
    file_name: str,
    url: str,
    cache_dir: Path,
    dataset_images_dir: Path,
    do_head_check: bool,
    verify_image: bool,
    link_mode: str,
    retries: int,
    timeout: int,
) -> Tuple[str, str, bool, str]:
    """
    Process one (file_name, url): write cache file, link into dataset.
    Returns: (file_name, url, ok, error_msg)
    """
    cache_path = url_to_cache_path(url, cache_dir)

    # Optional: HEAD check (quick fail-fast for 404 or non-image Content-Type)
    if do_head_check and not cache_path.exists():
        if not head_ok(url):
            return file_name, url, False, "HEAD check failed"

    # Download falls im Cache fehlend
    if not cache_path.exists():
        ok, msg = http_get_verified(
            url, cache_path, retries=retries, timeout=timeout, verify_image=verify_image
        )
        if not ok:
            return file_name, url, False, msg

    # Link/copy into dataset under the desired name
    dataset_img_path = dataset_images_dir / Path(file_name).name
    try:
        link_into_dataset(cache_path, dataset_img_path, link_mode=link_mode)
    except Exception as e:
        return file_name, url, False, f"link failed: {e}"

    return file_name, url, True, ""


# --------------------------------- Main -------------------------------------

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--manifest", required=True, type=Path, help="Path to train/val/test manifest .txt")
    ap.add_argument("--cache_dir", required=True, type=Path, help="e.g., data/cache/images")
    ap.add_argument("--dataset_images_dir", required=True, type=Path, help="ml/datasets/taco/images/<split>")
    ap.add_argument("--limit", type=int, default=0, help="optional: only warm-up N images (0=all)")
    ap.add_argument("--shuffle", action="store_true", help="randomize manifest order")
    ap.add_argument("--workers", type=int, default=min(8, max(2, (os.cpu_count() or 4) // 2)),
                    help="parallel downloads (default: half of CPU cores, max 8)")
    ap.add_argument("--retries", type=int, default=3, help="HTTP retries per file")
    ap.add_argument("--timeout", type=int, default=25, help="HTTP timeout seconds")
    ap.add_argument("--no-head-check", action="store_true", help="disable HEAD check before GET")
    ap.add_argument("--no-image-verify", action="store_true", help="disable PIL.verify() for received bytes")
    ap.add_argument("--link-mode", choices=["auto", "copy", "symlink"], default="auto",
                    help="dataset link behavior (default auto: Windows=copy, else=symlink)")
    args = ap.parse_args()

    pairs = load_manifest(args.manifest)
    if not pairs:
        print(f"[ERROR] Manifest empty or invalid: {args.manifest}", file=sys.stderr)
        sys.exit(1)

    if args.shuffle:
        random.shuffle(pairs)

    if args.limit and args.limit > 0:
        pairs = pairs[:args.limit]

    args.cache_dir.mkdir(parents=True, exist_ok=True)
    args.dataset_images_dir.mkdir(parents=True, exist_ok=True)

    do_head_check = not args.no_head_check
    verify_image = not args.no_image_verify

    # Parallel abarbeiten
    ok = fail = 0
    failures: List[Tuple[str, str, str]] = []  # (file_name, url, error)
    total = len(pairs)

    print(f"[INFO] prefetch start: total={total}, workers={args.workers}, "
          f"head_check={do_head_check}, img_verify={verify_image}, link_mode={args.link_mode}")

    with ThreadPoolExecutor(max_workers=max(1, args.workers)) as ex:
        futs = [
            ex.submit(
                process_one,
                file_name, url,
                args.cache_dir,
                args.dataset_images_dir,
                do_head_check,
                verify_image,
                args.link_mode,
                args.retries,
                args.timeout,
            )
            for (file_name, url) in pairs
        ]

        for i, fut in enumerate(as_completed(futs), 1):
            file_name, url, success, err = fut.result()
            if success:
                ok += 1
                if ok % 100 == 0 or ok == 1:
                    print(f"[OK] {ok}/{total}  {file_name}")
            else:
                fail += 1
                failures.append((file_name, url, err))
                print(f"[WARN] {file_name}  -> {err}", file=sys.stderr)

    # Missing report (CSV) next to the manifest
    if failures:
        miss_path = args.manifest.parent / f"missing_{args.manifest.name.replace('.txt','')}.csv"
        try:
            with miss_path.open("w", encoding="utf-8", newline="") as f:
                w = csv.writer(f)
                w.writerow(["file_name", "url", "error"])
                for fn, u, e in failures:
                    w.writerow([fn, u, e])
            print(f"[INFO] missing report: {miss_path} (fail={fail})")
        except Exception as e:
            print(f"[WARN] missing report write failed: {e}", file=sys.stderr)

    print(f"[INFO] prefetch finished. ok={ok}, fail={fail}")
    # Do not hard-fail even if there were errors (keep behavior)
    if fail > 0:
        sys.exit(0)


if __name__ == "__main__":
    main()