#!/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()