ml/scripts/ingest_feedback.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
ingest_feedback.py

Convert feedback/prediction logs into a YOLO-formatted dataset.
Reads:  feedback_logs/predictions.jsonl, feedback_logs/feedback.jsonl
Writes: ml/datasets/alami_user/{images,labels}/{train,val}, dataset.yaml, names.json

Usage:
  python ml/scripts/ingest_feedback.py --out ml/datasets/alami_user --split 0.85 --max 5000 --download
"""

from __future__ import annotations
import argparse
import json
import logging
import random
import shutil
import sys
from collections import OrderedDict
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import requests
from PIL import Image
import yaml


logger = logging.getLogger("ingest_feedback")


def make_session(retries: int = 3, backoff: float = 0.5) -> requests.Session:
    from urllib3.util.retry import Retry
    from requests.adapters import HTTPAdapter
    sess = requests.Session()
    retry = Retry(total=retries, read=retries, connect=retries, backoff_factor=backoff,
                  status_forcelist=(429, 500, 502, 503, 504), raise_on_status=False)
    adapter = HTTPAdapter(max_retries=retry)
    sess.mount("http://", adapter)
    sess.mount("https://", adapter)
    return sess


def download_image(url: str, out_dir: Path, session: Optional[requests.Session] = None, timeout: int = 30) -> Optional[Path]:
    sess = session or requests.Session()
    try:
        r = sess.get(url, stream=True, timeout=timeout)
        r.raise_for_status()
        base = Path((url or "").split("?")[0]).name or "img"
        stem = Path(base).stem[:60]
        suffix = Path(base).suffix or ".jpg"
        out_dir.mkdir(parents=True, exist_ok=True)
        tmp_path = out_dir / f"__tmp__{stem}{suffix}"
        with open(tmp_path, "wb") as f:
            shutil.copyfileobj(r.raw, f)
        return tmp_path
    except Exception as e:
        logger.warning("Download failed %s : %s", url, e)
        return None


def xyxy_to_yolo(xyxy: List[float], img_w: int, img_h: int) -> Tuple[float, float, float, float]:
    x1, y1, x2, y2 = xyxy[:4]
    x1 = max(0.0, min(float(x1), float(img_w)))
    x2 = max(0.0, min(float(x2), float(img_w)))
    y1 = max(0.0, min(float(y1), float(img_h)))
    y2 = max(0.0, min(float(y2), float(img_h)))
    w = max(1e-6, x2 - x1)
    h = max(1e-6, y2 - y1)
    cx = x1 + w / 2.0
    cy = y1 + h / 2.0
    cx, cy, w, h = cx / img_w, cy / img_h, w / img_w, h / img_h
    cx = min(max(cx, 0.0), 1.0)
    cy = min(max(cy, 0.0), 1.0)
    return (cx, cy, w, h)


def load_jsonl(path: Path) -> List[Dict]:
    if not path.exists():
        logger.info("Missing file: %s", path)
        return []
    out = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            s = line.strip()
            if not s:
                continue
            try:
                out.append(json.loads(s))
            except Exception as e:
                logger.warning("Invalid JSON in %s: %s", path, e)
    return out


def ensure_dirs(base: Path):
    for sub in ["images/train", "images/val", "labels/train", "labels/val"]:
        (base / sub).mkdir(parents=True, exist_ok=True)


def map_names_to_classes(names_path: Path) -> Dict[str, int]:
    if not names_path.exists():
        logger.warning("names.json not found at %s; will build class map from feedback.", names_path)
        return {}
    try:
        obj = json.loads(names_path.read_text(encoding="utf-8"))
        if isinstance(obj, dict):
            return {str(v): int(k) for k, v in obj.items()}
        if isinstance(obj, list):
            return {str(n): i for i, n in enumerate(obj)}
    except Exception:
        logger.exception("Failed to parse names.json at %s", names_path)
    return {}


def write_label_file(label_path: Path, class_id: int, yolo_box: Tuple[float, float, float, float]):
    label_path.write_text(
        f"{class_id} {yolo_box[0]:.6f} {yolo_box[1]:.6f} {yolo_box[2]:.6f} {yolo_box[3]:.6f}\n",
        encoding="utf-8"
    )


def main(argv=None):
    ap = argparse.ArgumentParser(description="Build YOLO dataset from feedback/prediction logs.")
    ap.add_argument("--predictions", default="feedback_logs/predictions.jsonl")
    ap.add_argument("--feedback", default="feedback_logs/feedback.jsonl")
    ap.add_argument("--out", default="ml/datasets/alami_user")
    ap.add_argument("--names", default="deploy/latest/names.json")
    ap.add_argument("--split", type=float, default=0.85)
    ap.add_argument("--max", type=int, default=5000)
    ap.add_argument("--download", action="store_true")
    ap.add_argument("--seed", type=int, default=1337)
    args = ap.parse_args(argv)

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")

    if not args.download:
        raise SystemExit("--download is required to create a usable training set.")

    out_base = Path(args.out)
    ensure_dirs(out_base)

    preds = load_jsonl(Path(args.predictions))
    fbs_all = load_jsonl(Path(args.feedback))
    logger.info("Loaded %d predictions, %d feedback entries", len(preds), len(fbs_all))

    fb_by_pid = OrderedDict()
    for fb in fbs_all:
        pid = fb.get("prediction_id")
        if pid:
            fb_by_pid[pid] = fb
    fbs = list(fb_by_pid.values())

    pred_by_id = {p.get("prediction_id"): p for p in preds if p.get("prediction_id")}
    class_map = map_names_to_classes(Path(args.names))

    pairs: List[Tuple[Dict, Dict, Dict]] = []
    for fb in fbs:
        pid = fb.get("prediction_id")
        if not pid:
            continue
        pitem = pred_by_id.get(pid)
        if not pitem:
            continue
        corr = fb.get("corrected_type")
        if not corr:
            continue
        plist = pitem.get("predictions") or []
        chosen = None
        for b in plist:
            if str(b.get("label")) == str(corr):
                chosen = b
                break
        if not chosen and plist:
            chosen = max(plist, key=lambda x: float(x.get("conf", 0.0)))
        if not chosen:
            continue
        pairs.append((pitem, fb, chosen))
        if len(pairs) >= args.max:
            break

    if not pairs:
        logger.warning("No matched feedback/predictions. Nothing to do.")
        return

    random.seed(args.seed)
    random.shuffle(pairs)

    seen_names = set(str(fb.get("corrected_type")) for _, fb, _ in pairs if fb.get("corrected_type"))
    if not class_map:
        class_map = {n: i for i, n in enumerate(sorted(seen_names))}
        logger.info("Class map generated from feedback: %s", class_map)

    sess = make_session()

    train_cut = int(len(pairs) * args.split)
    splits = [("train", pairs[:train_cut]), ("val", pairs[train_cut:])]

    for split_name, items in splits:
        imgs_dir = out_base / f"images/{split_name}"
        labels_dir = out_base / f"labels/{split_name}"
        for pred, fb, chosen in items:
            img_url = pred.get("image_url") or pred.get("path")
            if not img_url:
                continue
            tmp_path = download_image(img_url, imgs_dir, session=sess)
            if not tmp_path:
                continue
            try:
                with Image.open(tmp_path) as im:
                    img_w, img_h = im.size
                    im.verify()
            except Exception:
                tmp_path.unlink(missing_ok=True)
                continue

            pid = str(pred.get("prediction_id", ""))[:8]
            base = Path((img_url or "").split("?")[0]).name or "img"
            stem = Path(base).stem[:60]
            suffix = Path(base).suffix or ".jpg"
            final_name = f"{stem}__{pid}{suffix}"
            final_path = imgs_dir / final_name
            if final_path.exists():
                final_path.unlink()
            tmp_path.rename(final_path)

            xyxy = chosen.get("xyxy") or chosen.get("box")
            if not xyxy or len(xyxy) < 4:
                continue
            yolo_box = xyxy_to_yolo(xyxy, img_w, img_h)
            cx, cy, w, h = yolo_box
            if not (0 < w <= 1 and 0 < h <= 1):
                continue

            cls_name = str(fb.get("corrected_type"))
            if cls_name not in class_map:
                class_map[cls_name] = len(class_map)
            cls_id = int(class_map[cls_name])

            label_path = labels_dir / (final_path.stem + ".txt")
            write_label_file(label_path, cls_id, (cx, cy, w, h))

    names_arr = [None] * len(class_map)
    for k, v in class_map.items():
        if 0 <= v < len(names_arr):
            names_arr[v] = str(k)
    for i, v in enumerate(names_arr):
        if v is None:
            names_arr[i] = str(i)

    ds = {
        "path": str(out_base.resolve()),
        "train": "images/train",
        "val": "images/val",
        "test": "images/val",
        "nc": len(names_arr),
        "names": names_arr,
    }

    (out_base / "dataset.yaml").write_text(yaml.safe_dump(ds, sort_keys=False), encoding="utf-8")
    (out_base / "names.json").write_text(json.dumps(names_arr, ensure_ascii=False, indent=2), encoding="utf-8")

    n_train_img = len(list((out_base / "images/train").glob("*")))
    n_train_lab = len(list((out_base / "labels/train").glob("*.txt")))
    n_val_img = len(list((out_base / "images/val").glob("*")))
    n_val_lab = len(list((out_base / "labels/val").glob("*.txt")))
    logger.info("Done. Train images=%d labels=%d | Val images=%d labels=%d | Classes=%d",
                n_train_img, n_train_lab, n_val_img, n_val_lab, len(names_arr))


if __name__ == "__main__":
    main()