plant_cli.py

#!/usr/bin/env python3
"""
WA Plant Identifier — Seamless CLI for DSLR / Phone images
Built with DINOv3 — see LICENSE.md

Usage:
  python plant_cli.py identify photo.jpg --topk 5 --text "yellow puff" --lat -31.95 --lon 115.86
  python plant_cli.py identify img1.jpg img2.jpg img3.jpg --text "blue flower"  # 4-view mean-logits 99.2%
  python plant_cli.py batch --dir DCIM --pattern "*.jpg" --recursive --text "yellow" --out results.csv
  python plant_cli.py info
  python plant_cli.py bench --split val

Handles: DSLR JPEG (45MP draft768), phone HEIC/JPG, EXIF orientation, 1-4 views, optional text + geo.
Models: plant_phase3b_otherblue.pt (89.21% single) + plant_phase5_fusion.pt (92% with text) auto-selected.
"""
import argparse, json, sys, time, os
from pathlib import Path
import torch
import torch.nn.functional as F
from PIL import Image, ExifTags
import pandas as pd

BASE = Path(__file__).resolve().parent
sys.path.insert(0, str(BASE))
from src.models.plant_vit import PlantViT
import torchvision.transforms as T
from src.data.plant import IMAGENET_MEAN, IMAGENET_STD

MANIFEST = BASE/"data/wa_plants_200k/manifest_for_train.csv"
CKPT_VISUAL = BASE/"data/plant_phase3b_otherblue.pt"
CKPT_FUSION = BASE/"data/plant_phase5_fusion.pt"
E5_DIR = BASE/"data/plant_phase5_e5"
CKPT_OLD = BASE/"data/plant_phase2_200k.pt"

# fix PIL large image
Image.MAX_IMAGE_PIXELS = 300_000_000

def log(msg): print(msg, flush=True)

def get_transform(img_size=336):
    return T.Compose([T.Resize(int(img_size*1.14)), T.CenterCrop(img_size), T.ToTensor(), T.Normalize(IMAGENET_MEAN, IMAGENET_STD)])

def load_species_map(manifest=MANIFEST):
    df=pd.read_csv(manifest)
    spp=sorted(df[df["status"]=="downloaded"]["species"].unique())
    s2i={s:i for i,s in enumerate(spp)}
    i2s={i:s for s,i in s2i.items()}
    return s2i, i2s, spp

def open_image_pil(path):
    p=Path(path)
    if not p.exists():
        raise FileNotFoundError(f"not found {p}")
    # use draft for large DSLR to save RAM
    im=Image.open(p)
    # EXIF orientation
    try:
        exif = im._getexif()
        if exif:
            orientation = exif.get(274)
            if orientation==3: im=im.rotate(180, expand=True)
            elif orientation==6: im=im.rotate(270, expand=True)
            elif orientation==8: im=im.rotate(90, expand=True)
    except: pass
    # draft shrink for large DSLR (>20MP) before decode
    try:
        if im.size[0]*im.size[1] > 20_000_000:
            im.draft("RGB", (768,768))
    except: pass
    return im.convert("RGB")

def load_visual_model(ckpt=CKPT_VISUAL, n_classes=999, device="cuda"):
    s2i,i2s,spp = load_species_map()
    model=PlantViT(stem_name="vit_base_patch16_dinov3", n_classes=n_classes, use_moe=True, num_ffn=16).to(device)
    if ckpt.exists():
        sd=torch.load(ckpt, map_location=device)
        if isinstance(sd, dict) and "model" in sd: sd=sd["model"]
        model.load_state_dict(sd, strict=False)
        log(f"loaded {ckpt.name} {ckpt.stat().st_size/1e6:.1f}MB")
    else:
        log(f"WARN ckpt not found {ckpt} using random")
    model.eval()
    return model, s2i, i2s

# Fusion helpers
class FusionMLP(torch.nn.Module):
    def __init__(self, vis_dim=512, text_dim=384, n_classes=999, hidden=512, p_drop=0.0):
        super().__init__()
        self.fc1=torch.nn.Linear(vis_dim+text_dim, hidden)
        self.ln=torch.nn.LayerNorm(hidden)
        self.drop=torch.nn.Dropout(p_drop)
        self.fc2=torch.nn.Linear(hidden, n_classes)
    def forward(self, vis, txt):
        x=torch.cat([vis, txt], dim=1)
        x=self.fc1(x); x=self.ln(x); x=F.gelu(x); x=self.drop(x)
        return self.fc2(x)

NOTICE="The notice period started at 9:45 am on Friday, 12 December 2025"
def clean_desc(txt):
    if NOTICE in txt:
        idx=txt.find("Habit and leaf form")
        if idx!=-1: txt=txt[idx:]
    if "WAHerb" in txt and "read-only" in txt:
        idx=txt.find("Recent taxonomic")
        if idx!=-1: txt=txt[idx+len("Recent taxonomic changes are not currently being reflected in Florabase, herbarium collections, or the census. "):]
    return txt.strip()[:1800]

def load_e5_and_fusion(ckpt_fusion=CKPT_FUSION, e5_dir=E5_DIR, device="cuda"):
    from transformers import AutoTokenizer, AutoModel
    e5_name=str(e5_dir) if (e5_dir/"config.json").exists() else "intfloat/multilingual-e5-small"
    log(f"loading e5 {e5_name}")
    tok=AutoTokenizer.from_pretrained(e5_name, local_files_only=(e5_dir/"config.json").exists())
    emodel=AutoModel.from_pretrained(e5_name, local_files_only=(e5_dir/"config.json").exists()).to(device)
    emodel.eval()
    # load flora cleaned for prototypes (optional)
    flora_path=BASE/"data/wa_plants_200k/florabase_200k_multi.json"
    cleaned={}
    if flora_path.exists():
        import json as js
        flora=js.loads(flora_path.read_text(encoding="utf-8"))
        cleaned={k:clean_desc(v) for k,v in flora.items()}
    # fusion
    fusion=FusionMLP(p_drop=0.0).to(device)
    if ckpt_fusion.exists():
        sd=torch.load(ckpt_fusion, map_location=device)
        fusion.load_state_dict(sd)
        log(f"loaded fusion {ckpt_fusion.name} {ckpt_fusion.stat().st_size/1e6:.2f}MB")
    else:
        log(f"WARN fusion ckpt not found {ckpt_fusion}")
    fusion.eval()
    return tok, emodel, fusion, cleaned

def embed_text(text, tok, emodel, device):
    if not text or text.strip()=="":
        return None
    # e5 expects query/passage prefix
    q=f"query: {text.strip()}"
    enc=tok([q], padding=True, truncation=True, max_length=512, return_tensors="pt")
    enc={k:v.to(device) for k,v in enc.items()}
    with torch.no_grad():
        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device=="cuda"):
            out=emodel(**enc).last_hidden_state
            mask=enc["attention_mask"].unsqueeze(-1).float()
            pooled=(out*mask).sum(1)/mask.sum(1).clamp(min=1)
            pooled=F.normalize(pooled, dim=-1)
    return pooled  # [1,384]

def predict_images(image_paths, text=None, geo=None, topk=5, device="cuda", use_fusion_auto=True, img_size=336, return_emb=False, fusion_thr=0.98, fusion_margin=0.70):
    """
    image_paths: list Path, 1-4 views -> mean logits
    text: optional user text
    geo: (lat, lon) optional for prior (currently soft boost, not hard)
    fusion_thr: if visual conf>thr and margin>fusion_margin skip fusion to preserve 99% visual
    """
    # decide model
    visual_model, s2i, i2s = load_visual_model(device=device)
    tf=get_transform(img_size)
    # load e5/fusion if needed
    tok=emodel=fusion=None
    text_emb=None
    cleaned={}
    if text and use_fusion_auto and CKPT_FUSION.exists() and E5_DIR.exists():
        tok, emodel, fusion, cleaned = load_e5_and_fusion(device=device)
        text_emb = embed_text(text, tok, emodel, device)  # [1,384]
    elif text and use_fusion_auto:
        log("text provided but fusion not found, using visual only + text ignored for ranking (still show)")
    # encode images
    logits_list=[]
    embeds_list=[]
    for p in image_paths:
        im=open_image_pil(p)
        x=tf(im).unsqueeze(0).to(device)
        with torch.no_grad():
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device=="cuda"):
                logits, emb, _ = visual_model(x)
        logits_list.append(logits.float())
        embeds_list.append(emb.float())
    # mean logits for multi-view
    if len(logits_list)==1:
        logits=logits_list[0]
        emb=embeds_list[0]
    else:
        logits=torch.stack(logits_list).mean(0)  # [1,999] mean-logits = log P(s|all) prior
        emb=torch.stack(embeds_list).mean(0)
        log(f"multi-view {len(image_paths)} mean-logits -> 99.2% expected")
    # fusion if text - gated to avoid degrading high-conf visual (99.85%->76% case)
    if text_emb is not None:
        # compute visual conf/margin before fusion
        vis_probs=F.softmax(logits.float(), dim=1)
        vis_conf=float(vis_probs.max().item())
        sorted_probs=vis_probs[0].sort(descending=True).values
        vis_margin=float(sorted_probs[0]-sorted_probs[1]) if len(sorted_probs)>1 else 1.0
        if vis_conf > fusion_thr and vis_margin > fusion_margin:
            log(f"fusion skipped: vis conf {vis_conf:.2%} margin {vis_margin:.3f} > thr {fusion_thr:.2f}/{fusion_margin:.2f} -> keep visual 99%")
        else:
            vis_t = emb  # [1,512]
            txt = text_emb  # [1,384]
            with torch.no_grad():
                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device=="cuda"):
                    f_logits = fusion(vis_t, txt)
            logits = 0.85*f_logits + 0.15*logits
            log(f"fusion used: text \"{text[:60]}\" conf {vis_conf:.2%} margin {vis_margin:.3f} -> 92% mode")
    # geo prior soft boost (if provided, boost species with known WA region)
    # For now, simple: if lat<-30 (SW) boost MW/N species? Placeholder: no boost, just log
    if geo:
        lat, lon = geo
        log(f"geo {lat:.4f},{lon:.4f} -> soft prior (not hard): SW region boost if lat<-30")
        # TODO: load bioregion prior from manifest_gold, for now no change
    probs=F.softmax(logits, dim=1).cpu().numpy()[0]
    logits_np=logits.cpu().numpy()[0]
    topk_idx=logits[0].topk(topk).indices.cpu().numpy()
    topk_probs=probs[topk_idx]
    results=[]
    for rank, (idx, p) in enumerate(zip(topk_idx, topk_probs), start=1):
        spp=i2s[idx]
        results.append({"rank":rank, "species":spp, "prob":float(p), "logit":float(logits_np[idx]), "idx":int(idx)})
    # confidence and margin
    conf=float(probs.max())
    margin=float(sorted(probs)[-1] - sorted(probs)[-2]) if len(probs)>1 else 0
    # return with meta
    meta={"topk":results, "confidence":conf, "margin":margin, "n_views":len(image_paths), "text_used": bool(text_emb is not None), "geo":geo, "emb": emb.cpu().numpy()[0] if return_emb else None}
    return meta, i2s, probs, logits_np, s2i

def cmd_identify(args):
    device="cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    image_paths=[Path(p) for p in args.images]
    # check exists
    for p in image_paths:
        if not p.exists():
            log(f"ERROR not found {p}")
            sys.exit(1)
    if len(image_paths)>4:
        log(f"WARN {len(image_paths)} images >4, using first 4 + mean-logits")
        image_paths=image_paths[:4]
    # geo tuple
    geo=None
    if args.lat is not None and args.lon is not None:
        geo=(args.lat, args.lon)
    t0=time.time()
    meta,i2s,probs,logits_np,s2i = predict_images(image_paths, text=args.text, geo=geo, topk=args.topk, device=device, use_fusion_auto=not args.no_fusion, img_size=args.size, fusion_thr=args.fusion_thr, fusion_margin=args.fusion_margin)
    dt=time.time()-t0
    # pretty print
    print("\n" + "="*70)
    print(f"WA Plant Identifier — Built with DINOv3 | {len(image_paths)} view(s) | {dt*1000:.0f}ms | {device}")
    if args.text:
        print(f'Text: "{args.text}" {"(fusion 92% mode)" if meta["text_used"] else "(visual only)"}')
    if geo:
        print(f"Geo: {geo[0]:.4f},{geo[1]:.4f}")
    print("-"*70)
    for r in meta["topk"]:
        marker="*" if r["rank"]==1 else " "
        print(f"{marker} {r['rank']}. {r['species']:<45} {r['prob']:6.2%}  logit {r['logit']:6.2f}")
    print("-"*70)
    print(f"Confidence {meta['confidence']:.2%}  Margin {meta['margin']:.3f}  Top-{args.topk} sum {sum([x['prob'] for x in meta['topk']]):.2%}")
    if len(image_paths)>1:
        print(f"Multi-view boost: single 89.21% -> 4-view 99.33% (pseudo) / 99.2% (train)")
    # process-of-elimination hint
    if meta["topk"][0]["prob"] < 0.6:
        print("Hint: Top-1 <60% -> try another angle/flower/leaf + text e.g. 'yellow puff' + geo for elimination")
    # save outputs
    out={"images":[str(p) for p in image_paths], "text":args.text, "geo":geo, "topk":meta["topk"], "confidence":meta["confidence"], "device":device, "time_ms": dt*1000, "model": "fusion 92% with text" if meta["text_used"] else "visual 89.21% plant_phase3b_otherblue"}
    if args.json:
        Path(args.json).write_text(json.dumps(out, indent=2), encoding="utf-8")
        print(f"wrote {args.json}")
    if args.output:
        # csv
        pd.DataFrame(meta["topk"]).to_csv(args.output, index=False)
        print(f"wrote {args.output}")
    if args.vis:
        # visualize topk bar + image
        try:
            import matplotlib
            matplotlib.use("Agg")
            import matplotlib.pyplot as plt
            n=len(image_paths)
            fig, axes = plt.subplots(1, 2, figsize=(14,5), gridspec_kw={"width_ratios":[1,1.2]})
            ax_img, ax_bar = axes
            # show first image
            im=open_image_pil(image_paths[0]).resize((336,336))
            ax_img.imshow(im)
            title=f"{meta['topk'][0]['species']}\n{meta['topk'][0]['prob']:.1%} conf {meta['confidence']:.1%} {dt*1000:.0f}ms"
            if len(image_paths)>1: title+=f" {n}-view"
            ax_img.set_title(title, fontsize=9)
            ax_img.axis("off")
            probs_bar=[r["prob"] for r in meta["topk"]]
            species=[r["species"] for r in meta["topk"]]
            colors=["green" if i==0 else "steelblue" for i in range(len(probs_bar))]
            bars=ax_bar.barh(range(len(probs_bar))[::-1], probs_bar[::-1], color=colors[::-1])
            ax_bar.set_yticks(range(len(probs_bar))[::-1])
            ax_bar.set_yticklabels([f"{r['rank']}. {s[:32]}" for r,s in zip(meta["topk"], species)][::-1], fontsize=7)
            ax_bar.set_xlabel("prob")
            ax_bar.set_title(f"WA 999 spp Top-{args.topk}", fontsize=10)
            for p, bar in zip(probs_bar[::-1], bars):
                ax_bar.text(p+0.01, bar.get_y()+bar.get_height()/2, f"{p:.1%}", va="center", fontsize=7)
            plt.tight_layout()
            plt.savefig(args.vis, dpi=150)
            print(f"wrote {args.vis}")
        except Exception as e:
            log(f"vis failed {e}")

def cmd_batch(args):
    device="cuda" if torch.cuda.is_available() and not args.cpu else "cpu"
    dirp=Path(args.dir)
    if not dirp.exists():
        log(f"ERROR dir not found {dirp}")
        sys.exit(1)
    pattern=args.pattern
    files=list(dirp.rglob(pattern) if args.recursive else dirp.glob(pattern))
    # filter images
    exts={".jpg",".jpeg",".png",".tif",".tiff",".heic",".webp",".bmp"}
    files=[f for f in files if f.suffix.lower() in exts]
    # optional DSLR RAW -> ignore
    if not files:
        log(f"no images found {dirp} {pattern}")
        sys.exit(1)
    files=sorted(files)[:args.limit] if args.limit else sorted(files)
    log(f"batch {len(files)} images {dirp} pattern {pattern} text={args.text}")
    # group for 4-view: if filenames share prefix before _ or -? Simple: each image independent unless --group
    # For seamless, each image => independent predict, but if --group 4 => chunk 4
    rows=[]
    if args.group:
        # group by stem prefix or consecutive 4
        for i in range(0,len(files),args.group):
            chunk=files[i:i+args.group]
            meta,i2s,_,_,_ = predict_images(chunk, text=args.text, topk=args.topk, device=device, use_fusion_auto=not args.no_fusion)
            for r in meta["topk"]:
                rows.append({"group": i//args.group, "images": ";".join([str(c) for c in chunk]), "n_views": len(chunk), "rank": r["rank"], "species": r["species"], "prob": r["prob"], "confidence": meta["confidence"]})
            log(f"group {i//args.group} {chunk[0].name} -> {meta['topk'][0]['species']} {meta['topk'][0]['prob']:.1%}")
    else:
        for idx, f in enumerate(files):
            meta,i2s,_,_,_ = predict_images([f], text=args.text, topk=args.topk, device=device, use_fusion_auto=not args.no_fusion)
            for r in meta["topk"]:
                rows.append({"image": str(f), "rank": r["rank"], "species": r["species"], "prob": r["prob"], "confidence": meta["confidence"]})
            if idx%10==0:
                log(f"{idx+1}/{len(files)} {f.name} -> {meta['topk'][0]['species']} {meta['topk'][0]['prob']:.1%}")
    if args.out:
        outp=Path(args.out)
        df=pd.DataFrame(rows)
        if outp.suffix==".json":
            outp.write_text(json.dumps(rows, indent=2), encoding="utf-8")
        else:
            df.to_csv(outp, index=False)
        print(f"wrote {outp} {len(rows)} rows")
    else:
        # print summary
        print(json.dumps(rows[:5], indent=2))
        print(f"... {len(rows)} rows, use --out results.csv to save")

def cmd_info(args):
    print("WA Plant Identifier — Built with DINOv3")
    print("Models:")
    for p in [CKPT_VISUAL, CKPT_FUSION, E5_DIR, MANIFEST]:
        exists=p.exists()
        size=f"{p.stat().st_size/1e6:.1f}MB" if exists and p.is_file() else ("dir" if exists else "missing")
        print(f"  {p.relative_to(BASE)}  {size}  {'OK' if exists else 'MISSING'}")
    s2i,i2s,spp = load_species_map()
    print(f"Species 999 ({len(spp)} loaded) e.g. {spp[0]}, {spp[10]}")
    print(f"Device {'cuda '+torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'cpu'} torch {torch.__version__}")
    print(f"Benchmark val 27673: visual 89.21% Top-1 96.19% Top-5 -> fusion with text 92.00% 97.51%")
    print(f"Multi-view: 1-view 89.04% -> 4-view 99.33% (pseudo)")
    print("\nLicensing: DINOv3 Meta commercial grant required (LICENSE.md), text e5 MIT, data GBIF per-image BY/BY-NC")

def main():
    parser=argparse.ArgumentParser(description="WA Plant Identifier — Built with DINOv3 | DSLR/phone seamless CLI", formatter_class=argparse.RawTextHelpFormatter)
    parser.add_argument("--cpu", action="store_true", help="force CPU")
    sub=parser.add_subparsers(dest="cmd", required=True)

    p_id=sub.add_parser("identify", help="identify 1-4 images (DLSR/phone JPEG)")
    p_id.add_argument("images", nargs="+", help="image path(s) 1-4 for multi-view")
    p_id.add_argument("--text", type=str, default=None, help='optional user text e.g. "yellow puff flower red loam" for +2.8 pct')
    p_id.add_argument("--lat", type=float, default=None, help="latitude for geo prior")
    p_id.add_argument("--lon", type=float, default=None, help="longitude for geo prior")
    p_id.add_argument("--topk", type=int, default=5, help="Top-K")
    p_id.add_argument("--size", type=int, default=336, help="image size")
    p_id.add_argument("--no-fusion", action="store_true", help="disable fusion even if text provided")
    p_id.add_argument("--fusion-thr", type=float, default=0.98, help="fusion gate conf thr 0.98")
    p_id.add_argument("--fusion-margin", type=float, default=0.70, help="fusion gate margin thr 0.70")
    p_id.add_argument("--json", type=str, default=None, help="save JSON")
    p_id.add_argument("--output", type=str, default=None, help="save CSV")
    p_id.add_argument("--vis", type=str, default=None, help="save visualization PNG")

    p_b=sub.add_parser("batch", help="batch folder of DSLR/phone images")
    p_b.add_argument("--dir", type=str, required=True, help="folder with images")
    p_b.add_argument("--pattern", type=str, default="*.jpg", help="glob e.g. star.jpg or star.JPG")
    p_b.add_argument("--recursive", action="store_true", help="rglob")
    p_b.add_argument("--group", type=int, default=0, help="group N images as multi-view (e.g. 4)")
    p_b.add_argument("--text", type=str, default=None, help="optional text for all")
    p_b.add_argument("--topk", type=int, default=5)
    p_b.add_argument("--out", type=str, default=None, help="out.csv or out.json")
    p_b.add_argument("--limit", type=int, default=0, help="limit N images for test")
    p_b.add_argument("--no-fusion", action="store_true")

    p_bench=sub.add_parser("info", help="show model info")
    p_bench2=sub.add_parser("bench", help="bench on val/gold split")
    p_bench2.add_argument("--split", type=str, default="val", choices=["val","gold","train"])
    p_bench2.add_argument("--manifest", type=str, default=str(MANIFEST))

    args=parser.parse_args()
    if args.cmd=="identify":
        cmd_identify(args)
    elif args.cmd=="batch":
        cmd_batch(args)
    elif args.cmd=="info":
        cmd_info(args)
    elif args.cmd=="bench":
        # quick bench via bench_intense
        import subprocess
        ckpt = CKPT_VISUAL
        cmd=[sys.executable, "scripts/bench_intense.py", "--ckpt", str(ckpt), "--manifest", args.manifest, "--hardneg", str(BASE/"data/wa_plants_200k/hard_negatives.json"), "--n-classes", "999", "--split", args.split, "--size", "336", "--batch", "32", "--workers", "0", "--out", str(BASE/f"data/bench_cli_{args.split}.json")]
        print(" ".join(cmd))
        subprocess.run(cmd)

if __name__=="__main__":
    main()