Spaces:
Sleeping
Sleeping
| """Explicit adapters for additional paper methods; no automatic disease votes. | |
| SAM2 follows the author's training architecture, not the incompatible SAM1 | |
| inference_clr.py. Automatic points come from DeepLab, never a reference mask. | |
| SCOLD architecture adapted from enalis/scold, CC BY 4.0, Nguyen Quoc et al. | |
| """ | |
| import hashlib | |
| import json | |
| import os | |
| from pathlib import Path | |
| import shutil | |
| import subprocess | |
| import tempfile | |
| import tarfile | |
| import time | |
| import urllib.error | |
| import urllib.request | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| import torch | |
| from torch import nn | |
| SAM2_FILE = "01_models/03_trained_models/severity_rust.pt" | |
| SAM3_REPO = "facebook/sam3" | |
| SCOLD_REPO = "enalis/scold" | |
| ROBERTA_REPO = "FacebookAI/roberta-base" | |
| REVISIONS = { | |
| ROBERTA_REPO: "e2da8e2f811d1448a5b465c236feacd80ffbac7b", | |
| SAM3_REPO: "3c879f39826c281e95690f02c7821c4de09afae7", | |
| SCOLD_REPO: "f0a65acc70c028510e0fd87ce9295ccfee1c97b3", | |
| } | |
| IJ_VERSION = "1.54p" | |
| IJ_SHA256 = "2e1a09961dfb41cee66ddc821b2577a41a072566ce45a49bae69267099741e20" | |
| def download_native_source(url, path): | |
| """Retry transient upstream failures; checksum validation stays with callers.""" | |
| for attempt in range(3): | |
| try: | |
| with urllib.request.urlopen(url, timeout=60) as source: | |
| path.write_bytes(source.read()) | |
| return | |
| except urllib.error.HTTPError as exc: | |
| if exc.code not in {429, 500, 502, 503, 504} or attempt == 2: | |
| raise | |
| time.sleep(2 ** (attempt + 1)) | |
| def prepare_native_runtime(): | |
| """Install isolated R dependencies on Linux once; never alter system R.""" | |
| if os.name == "nt" or os.getenv("AGROMIND_INSTALL_NATIVE", "1") != "1": | |
| return | |
| folder = Path.home() / ".cache" / "agromind" / "native" | |
| folder.mkdir(parents=True, exist_ok=True) | |
| prefix, mamba = folder / "r-runtime", folder / "micromamba" | |
| marker = prefix / "agromind-pliman-2.1.0.ready" | |
| if not marker.exists(): | |
| if not mamba.exists(): | |
| archive = folder / "micromamba-2.9.0.tar.bz2" | |
| download_native_source("https://conda.anaconda.org/conda-forge/linux-64/micromamba-2.9.0-0.tar.bz2", archive) | |
| if hashlib.sha256(archive.read_bytes()).hexdigest() != "8761c382127e6363bd9e0a2451aa3ef90d071a79133f736e2f759a3bf13040dd": | |
| raise ValueError("Native runtime installer checksum mismatch") | |
| with tarfile.open(archive) as source: | |
| mamba.write_bytes(source.extractfile("bin/micromamba").read()) | |
| mamba.chmod(0o755) | |
| if hashlib.sha256(mamba.read_bytes()).hexdigest() != "366cd9cd8be14df1ab8ed50352a82111082a36686b2d389fdb79a92c3fafb3e3": | |
| raise ValueError("Native runtime executable checksum mismatch") | |
| env = {**os.environ, "MAMBA_ROOT_PREFIX": str(folder / "mamba-cache")} | |
| args = [str(mamba), "create", "-y", "-p", str(prefix), "-c", "conda-forge", "-c", "bioconda", | |
| "--strict-channel-priority", "r-base=4.4", "bioconductor-ebimage=4.48.0", "r-sf", "r-stars", "r-terra", | |
| "r-doparallel", "r-foreach", "r-rcpparmadillo", "cxx-compiler", "make"] | |
| log = folder / "setup.log" | |
| package = folder / "pliman_2.1.0.tar.gz" | |
| download_native_source("https://cloud.r-project.org/src/contrib/Archive/pliman/pliman_2.1.0.tar.gz", package) | |
| if hashlib.sha256(package.read_bytes()).hexdigest() != "e73b2aeb56e97ee54a4b0b73840d6f70c65bfd96301cd07c4064a2b3bc3afb5d": | |
| raise ValueError("pliman source checksum mismatch") | |
| with log.open("w", encoding="utf-8") as output: | |
| subprocess.run(args, env=env, stdout=output, stderr=subprocess.STDOUT, check=True, timeout=900) | |
| subprocess.run([str(mamba), "run", "-p", str(prefix), "R", "CMD", "INSTALL", str(package)], | |
| env=env, stdout=output, stderr=subprocess.STDOUT, check=True, timeout=600) | |
| subprocess.run([str(prefix / "bin" / "Rscript"), "-e", "library(pliman); library(EBImage)"], check=True, timeout=60) | |
| marker.write_text("pliman 2.1.0; EBImage 4.48.0; source checksums verified\n") | |
| os.environ["AGROMIND_RSCRIPT"] = str(prefix / "bin" / "Rscript") | |
| def seed_points(mask, limit=5): | |
| """One interior point in each largest component; never read expert labels.""" | |
| count, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), 8) | |
| points = [] | |
| for i in sorted(range(1, count), key=lambda i: -stats[i, cv2.CC_STAT_AREA])[:limit]: | |
| distance = cv2.distanceTransform((labels == i).astype(np.uint8), cv2.DIST_L2, 5) | |
| y, x = np.unravel_index(distance.argmax(), distance.shape) | |
| points.append([int(x), int(y)]) | |
| return points | |
| def parse_points(text, size): | |
| values = json.loads(text) | |
| if not isinstance(values, list) or not 1 <= len(values) <= 5: | |
| raise ValueError("Provide 1–5 symptom points as [[x,y], ...] in crop pixels.") | |
| a = np.asarray(values, dtype=float) | |
| if a.shape != (len(values), 2) or not np.isfinite(a).all() or (a < 0).any() or (a >= size).any(): | |
| raise ValueError("Every symptom point must lie inside the selected crop.") | |
| return a.astype(np.float32) | |
| class ImageEncoder(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| import timm | |
| self.swin = timm.create_model("swin_base_patch4_window7_224.ms_in22k", pretrained=False, features_only=True) | |
| self.fc1 = nn.Linear(1024 * 7 * 7, 512) | |
| def forward(self, x): | |
| return self.fc1(self.swin(x)[-1].flatten(1)) | |
| class RobertaEncoder(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| from transformers import RobertaModel | |
| self.roberta = RobertaModel(config) | |
| self.projection = nn.Linear(config.hidden_size, 512) | |
| def forward(self, tokens): | |
| x = self.roberta(**tokens).last_hidden_state | |
| # Author uses an unmasked mean. Process prompts individually so batch | |
| # padding cannot change embeddings compared with the published example. | |
| return self.projection(x[:, 0] + x.mean(dim=1)) | |
| class Scold(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| self.image_encoder = ImageEncoder() | |
| self.text_encoder = RobertaEncoder(config) | |
| self.t_prime = nn.Parameter(torch.tensor(np.log(.07), dtype=torch.float32)) | |
| self.b = nn.Parameter(torch.zeros(())) | |
| def load(kind, device, download, revision, provenance, clr_repo): | |
| if kind == "sam2": | |
| from sam2.build_sam import build_sam2 | |
| from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| from unittest.mock import patch | |
| # SAM2 eagerly moves RoPE caches to CUDA, then deep-copies them when | |
| # cloning attention layers. ZeroGPU cannot deepcopy its virtual CUDA | |
| # storage during startup. Construct entirely on CPU, then move weights. | |
| # This short scope runs during serialized startup, never during requests. | |
| with patch("torch.cuda.is_available", return_value=False), torch.device("cpu"): | |
| model = build_sam2("configs/sam2.1/sam2.1_hiera_t.yaml", None, device="cpu", apply_postprocessing=False) | |
| model.load_state_dict(torch.load(download(clr_repo, SAM2_FILE), map_location="cpu", weights_only=True), strict=True) | |
| model.eval().to(device) | |
| meta = provenance(clr_repo, SAM2_FILE) | |
| meta.update(architecture="SAM2.1 Hiera tiny", prompt_contract="Independent positive-point prompts, candidate 0 as in training; union for display", | |
| adaptation="Automatic points, if selected, come from DeepLab. Training used reference-derived points; those are prohibited in benchmark inference.") | |
| return model, SAM2ImagePredictor(model), meta | |
| if kind == "sam3": | |
| from transformers import Sam3Model, Sam3Processor | |
| dtype_name = os.getenv("AGROMIND_SAM3_DTYPE", "bfloat16" if device == "cuda" else "float32") | |
| if dtype_name not in {"float32", "bfloat16"}: | |
| raise ValueError("AGROMIND_SAM3_DTYPE must be float32 or bfloat16") | |
| model = Sam3Model.from_pretrained(SAM3_REPO, revision=revision(SAM3_REPO), | |
| use_safetensors=True, trust_remote_code=False, | |
| dtype=getattr(torch, dtype_name)).eval().to(device) | |
| return model, Sam3Processor.from_pretrained(SAM3_REPO, revision=revision(SAM3_REPO)), {**provenance(SAM3_REPO), "dtype": dtype_name} | |
| if kind == "scold": | |
| from transformers import RobertaConfig, RobertaTokenizer | |
| rev = revision(ROBERTA_REPO) | |
| config = RobertaConfig.from_pretrained(ROBERTA_REPO, revision=rev) | |
| model = Scold(config) | |
| model.load_state_dict(torch.load(download(SCOLD_REPO, "scold.pth"), map_location="cpu", weights_only=True), strict=True) | |
| meta = provenance(SCOLD_REPO, "scold.pth") | |
| meta.update(tokenizer_revision=rev, license="CC BY 4.0", credit="Khang Nguyen Quoc, Lan Le Thi Thu, Luyl-Da Quach", | |
| source="https://huggingface.co/enalis/scold", adaptation="Trusted local architecture, strict numeric weight loading; no pretrained-backbone substitution") | |
| return model.eval().to(device), RobertaTokenizer.from_pretrained(ROBERTA_REPO, revision=rev), meta | |
| raise ValueError(kind) | |
| def predict(kind, image, model, proc, meta, device, prompts=None, points=None, text="orange spot"): | |
| if kind == "sam2": | |
| if points is None or len(points) == 0: | |
| raise ValueError("SAM2 needs symptom points. No seed is unavailable evidence, not zero rust.") | |
| proc.set_image(np.asarray(image).copy()) | |
| masks, scores, _ = proc.predict(point_coords=np.asarray(points, np.float32)[:, None, :], | |
| point_labels=np.ones((len(points), 1), np.int32), multimask_output=True) | |
| masks = np.asarray(masks) | |
| # SAM2 squeezes a single prompt batch; preserve independent prompts. | |
| candidates = masks[None] if masks.ndim == 3 else masks | |
| mask = np.any(candidates[:, 0].astype(bool), axis=0) | |
| if mask.shape != (image.height, image.width): | |
| raise ValueError("SAM2 returned a different pixel grid.") | |
| return mask, {**meta, "points_xy": np.asarray(points).tolist(), "mask_index": 0, | |
| "predicted_mask_quality": np.asarray(scores).reshape(-1, 3)[:, 0].tolist(), "not_disease_probability": True} | |
| if kind == "sam3": | |
| inputs = proc(images=image, text=text, return_tensors="pt").to(device, dtype=model.dtype) | |
| output = model(**inputs) | |
| result = proc.post_process_instance_segmentation(output, threshold=.5, mask_threshold=.5, | |
| target_sizes=inputs["original_sizes"].tolist())[0] | |
| masks = result["masks"].detach().cpu().numpy() | |
| mask = np.any(masks > 0, axis=0) if len(masks) else np.zeros((image.height, image.width), bool) | |
| return mask, {**meta, "text_prompt": text, "instance_count": len(masks), "score_threshold": .5, | |
| "mask_threshold": .5, "warning": "Generic text-prompt segmentation, not a rust diagnosis. Zero masks does not mean healthy."} | |
| if kind == "scold": | |
| a = np.asarray(image.resize((224, 224), Image.Resampling.BILINEAR)).astype(np.float32) / 255 | |
| x = torch.from_numpy(a.transpose(2, 0, 1)).unsqueeze(0).to(device) | |
| image_emb = nn.functional.normalize(model.image_encoder(x), dim=-1) | |
| ranking = [] | |
| for category, text in prompts: | |
| tokens = proc(text, return_tensors="pt", truncation=True).to(device) | |
| text_emb = nn.functional.normalize(model.text_encoder(tokens), dim=-1) | |
| ranking.append({"category": category, "prompt": text, "cosine_similarity": float((image_emb @ text_emb.T).item())}) | |
| return {"ranking": sorted(ranking, key=lambda r: -r["cosine_similarity"]), "provenance": meta, | |
| "warning": "SCOLD symptom similarity is not a diagnostic vote or calibrated probability."}, None | |
| raise ValueError(kind) | |
| def imagej_jar(): | |
| folder = Path(os.getenv("AGROMIND_TOOL_CACHE", str(Path.home() / ".cache" / "agromind"))) | |
| folder.mkdir(parents=True, exist_ok=True) | |
| jar = folder / f"ij-{IJ_VERSION}.jar" | |
| if not jar.exists(): | |
| with urllib.request.urlopen(f"https://repo.maven.apache.org/maven2/net/imagej/ij/{IJ_VERSION}/ij-{IJ_VERSION}.jar", timeout=60) as source: | |
| data = source.read(5_000_000) | |
| if hashlib.sha256(data).hexdigest() != IJ_SHA256: | |
| raise ValueError("ImageJ download checksum differs from the reviewed artifact.") | |
| jar.write_bytes(data) | |
| if hashlib.sha256(jar.read_bytes()).hexdigest() != IJ_SHA256: | |
| raise ValueError("ImageJ cached JAR checksum mismatch.") | |
| return jar | |
| def classical(kind, image, leaf, hsb="0,50,60,255,20,255", index="NGRDI"): | |
| """Invoke actual ImageJ/pliman runtimes, never relabel a Python approximation.""" | |
| base = Path(__file__).resolve().parent | |
| with tempfile.TemporaryDirectory(prefix="agromind_method_") as tmp: | |
| folder = Path(tmp) | |
| image.save(folder / "image.png") | |
| Image.fromarray(leaf.astype(np.uint8) * 255).save(folder / "leaf.png") | |
| if kind == "imagej": | |
| values = [int(v) for v in hsb.split(",")] | |
| if len(values) != 6 or any(not 0 <= v <= 255 for v in values) or any(values[i] > values[i+1] for i in (0, 2, 4)): | |
| raise ValueError("HSB needs six integers in 0–255: min,max for hue, saturation, brightness.") | |
| java = shutil.which("java") | |
| if not java: | |
| raise RuntimeError("ImageJ requires Java 11+; runtime is unavailable.") | |
| args = [java, "-Djava.awt.headless=true", "--class-path", str(imagej_jar()), str(base / "ImageJMask.java"), | |
| str(folder / "image.png"), str(folder / "mask.png"), *map(str, values)] | |
| meta = {"library": "ImageJ", "version": IJ_VERSION, "jar_sha256": IJ_SHA256, "hsb_ranges": values} | |
| elif kind == "pliman": | |
| if index not in {"R", "G", "B", "NGRDI"}: | |
| raise ValueError("Unsupported pliman index.") | |
| rscript = os.getenv("AGROMIND_RSCRIPT") or shutil.which("Rscript") | |
| if not rscript: | |
| raise RuntimeError("pliman requires R and its installed package; runtime is unavailable.") | |
| args = [rscript, str(base / "pliman_mask.R"), str(folder / "image.png"), str(folder / "leaf.png"), str(folder / "mask.png"), index] | |
| meta = {"library": "pliman", "index": index, "threshold": "Otsu"} | |
| else: | |
| raise ValueError(kind) | |
| process = subprocess.run(args, capture_output=True, text=True, timeout=120, check=False) | |
| if process.returncode or not (folder / "mask.png").exists(): | |
| raise RuntimeError(f"{kind} failed: {process.stderr[-2000:]}") | |
| with Image.open(folder / "mask.png") as im: | |
| mask = np.asarray(im.convert("L")) > 0 | |
| if mask.shape != leaf.shape: | |
| raise ValueError("Classical method returned a different pixel grid.") | |
| meta.update(runtime_output=process.stdout[-1000:], | |
| warning="Configurable colour baseline, not disease recognition. Original paper threshold parameters were not published; this is an explicit adaptation.") | |
| return mask & leaf, meta | |