scripts/merge_ext_train.py

#!/usr/bin/env python3
"""Merge all normalized single-class weed datasets, dedupe, split, train.

Gathers every /opt/weeds/ext/*/yolo (images/ + labels/), dedupes by image
content hash across sources (the grass-weeds RF100 set is re-hosted several
times), splits 80/10/10, and trains yolo11n + yolo11s at 640 for
"broadleaf weed vs grass" (single class). Evaluates on the held-out split.
"""
import glob
import hashlib
import json
import random
import shutil
from pathlib import Path
from PIL import Image
from ultralytics import YOLO

EXT = Path("/opt/weeds/ext")
DS = Path("/opt/weeds/broadleaf_dataset")
RUNS = "/opt/weeds/broadleaf_runs"
IMG_EXT = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

# The dock-in-grass "grass-weeds" set is re-hosted 6× (RF100 / RF100-VL / HF
# Francesco / HF LibreYOLO / Kaggle jaidalmotra / cotton-weed) with different
# augmentations that defeat perceptual dedup. Including several copies would
# overtrain + leak the same scene across train/test. So keep exactly ONE clean
# grass-weeds source (rf100vl — densest annotations) plus the genuinely
# distinct datasets. dhash still dedups within these.
INCLUDE = {
    "rf_rf100vl_grassweeds",   # dock-in-grass, ground-level — CORE deployment match
    "rf_augstartups_weeds",    # distinct generic weed set (diversity)
    "rf_weedswf1tx",           # aerial drone weeds (diversity)
    "rf_weedsffm3d",           # nettle/thistle (diversity)
}
SOURCE_PRIORITY = [
    "rf_rf100vl_grassweeds", "rf_augstartups_weeds",
    "rf_weedswf1tx", "rf_weedsffm3d",
]


def md5(path, buf=1 << 16):
    h = hashlib.md5()
    with open(path, "rb") as f:
        while (b := f.read(buf)):
            h.update(b)
    return h.hexdigest()


def dhash(path, size=16):
    """Perceptual hash: catches the same image across different JPEG encodings."""
    try:
        im = Image.open(path).convert("L").resize((size + 1, size), Image.BILINEAR)
    except Exception:
        return None
    px = list(im.getdata())
    bits = 0
    for r in range(size):
        row = px[r * (size + 1):(r + 1) * (size + 1)]
        for c in range(size):
            bits = (bits << 1) | (1 if row[c] < row[c + 1] else 0)
    return bits


def find_image(images_dir, stem):
    for e in IMG_EXT:
        p = images_dir / f"{stem}{e}"
        if p.exists():
            return p
    hits = list(images_dir.glob(f"{stem}.*"))
    return hits[0] if hits else None


def gather():
    records = []          # (src, img_path, label_path)
    per_src = {}
    found = {p.split("/")[-2]: p for p in glob.glob(str(EXT / "*" / "yolo"))}
    found = {s: p for s, p in found.items() if s in INCLUDE}
    ordered = [s for s in SOURCE_PRIORITY if s in found] + \
              [s for s in found if s not in SOURCE_PRIORITY]
    for src in ordered:
        yolo = Path(found[src])
        idir, ldir = yolo / "images", yolo / "labels"
        if not idir.exists() or not ldir.exists():
            continue
        n = 0
        for lp in ldir.glob("*.txt"):
            ip = find_image(idir, lp.stem)
            if ip is None:
                continue
            records.append((src, ip, lp))
            n += 1
        # also images without labels = background negatives
        labeled = {lp.stem for lp in ldir.glob("*.txt")}
        for ip in idir.iterdir():
            if ip.suffix.lower() in IMG_EXT and ip.stem not in labeled:
                records.append((src, ip, None))
                n += 1
        per_src[src] = n
    return records, per_src


def normalize_label_text(lp):
    if lp is None:
        return ""
    out = []
    for ln in Path(lp).read_text().splitlines():
        p = ln.split()
        if len(p) >= 5:
            out.append("0 " + " ".join(p[1:5]))
    return "\n".join(out)


def build():
    records, per_src = gather()
    print("per-source raw:", json.dumps(per_src))
    # dedupe by exact md5 AND perceptual dhash (catches re-encoded copies of
    # the same grass-weeds images across HF/Roboflow/Kaggle re-hosts)
    seen_md5, seen_dhash = set(), set()
    deduped = []
    dups = 0
    kept_by_src = {}
    for src, ip, lp in records:
        m = None
        try:
            m = md5(ip)
        except Exception:
            continue
        dh = dhash(ip)
        if m in seen_md5 or (dh is not None and dh in seen_dhash):
            dups += 1
            continue
        seen_md5.add(m)
        if dh is not None:
            seen_dhash.add(dh)
        deduped.append((src, ip, lp))
        kept_by_src[src] = kept_by_src.get(src, 0) + 1
    print(f"total {len(records)} → deduped {len(deduped)} ({dups} dup images dropped)")
    print("kept per source:", json.dumps(kept_by_src))

    rng = random.Random(42)
    rng.shuffle(deduped)
    n = len(deduped)
    n_val, n_test = int(n * 0.1), int(n * 0.1)
    split_of = {}
    for i, rec in enumerate(deduped):
        split_of[i] = "val" if i < n_val else "test" if i < n_val + n_test else "train"

    counts = {"train": 0, "val": 0, "test": 0}
    box_counts = {"train": 0, "val": 0, "test": 0}
    for i, (src, ip, lp) in enumerate(deduped):
        split = split_of[i]
        idir = DS / "images" / split
        ldir = DS / "labels" / split
        idir.mkdir(parents=True, exist_ok=True)
        ldir.mkdir(parents=True, exist_ok=True)
        stem = f"{src}__{ip.stem}"
        di = idir / f"{stem}{ip.suffix.lower()}"
        if not di.exists():
            di.symlink_to(ip.resolve())
        txt = normalize_label_text(lp)
        (ldir / f"{stem}.txt").write_text(txt + ("\n" if txt else ""))
        counts[split] += 1
        box_counts[split] += len(txt.splitlines()) if txt else 0
    (DS / "data.yaml").write_text(
        f"path: {DS}\ntrain: images/train\nval: images/val\ntest: images/test\n"
        "nc: 1\nnames:\n  0: weed\n")
    print("split images:", counts, "boxes:", box_counts)
    return counts


def train_eval():
    data = str(DS / "data.yaml")
    for name, weights, batch in [
        ("broadleaf-yolo11n-640", "/opt/rumex/weights/yolo11n.pt", 64),
        ("broadleaf-yolo11s-640", "/opt/rumex/weights/yolo11s.pt", 32),
    ]:
        best = Path(RUNS) / name / "weights" / "best.pt"
        if not best.exists():
            print(f"=== TRAIN {name}", flush=True)
            YOLO(weights).train(
                data=data, epochs=100, imgsz=640, batch=batch, patience=20,
                device=0, workers=8, project=RUNS, name=name, exist_ok=True,
                seed=42, plots=True)
        m = YOLO(str(best)).val(data=data, split="test", imgsz=640, device=0,
                                project=RUNS, name=f"{name}-test", exist_ok=True)
        rep = {"model": name, "test_mAP50": float(m.box.map50),
               "test_mAP50_95": float(m.box.map),
               "test_precision": float(m.box.mp), "test_recall": float(m.box.mr)}
        (Path(RUNS) / name / "test_eval.json").write_text(json.dumps(rep, indent=2))
        print(f"=== {name} TEST: {json.dumps(rep)}", flush=True)
    print("BROADLEAF_TRAIN_DONE", flush=True)


if __name__ == "__main__":
    build()
    train_eval()