Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| """DermaFusion inference engine — live embedding (SigLIP2 + DermLIP) -> late-fusion LightGBM -> 6-class | |
| skin-disease screening/classification. CPU / MPS / CUDA. Loads once; .predict(PIL.Image) -> structured result. | |
| NOT a diagnosis; clinical decision-support screening aid only. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| # Prevent the macOS libomp double-load segfault when transformers + open_clip co-load (must precede torch). | |
| os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") | |
| os.environ.setdefault("OMP_NUM_THREADS", "4") | |
| os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") | |
| import json, math | |
| from pathlib import Path | |
| import numpy as np, torch | |
| from PIL import Image, ImageOps | |
| CLASS_LABELS = { # human-readable + malignancy | |
| "bcc": ("Basal cell carcinoma", "malignant"), | |
| "scc": ("Squamous cell carcinoma", "malignant"), | |
| "ak_ack": ("Actinic keratosis", "malignant"), # pre-malignant, grouped high-risk | |
| "mel_mm": ("Melanoma", "malignant"), | |
| "nevus": ("Melanocytic nevus", "benign"), | |
| "seborrheic_keratosis": ("Seborrheic keratosis", "benign"), | |
| } | |
| def _pick_device(dev=None): | |
| # CPU is the tested-stable default (co-loading transformers SigLIP2 + open_clip DermLIP can segfault on | |
| # MPS). CUDA is safe when present. Callers may force a device explicitly. | |
| if dev: return torch.device(dev) | |
| if torch.cuda.is_available(): return torch.device("cuda") | |
| return torch.device("cpu") | |
| class DermaFusionEngine: | |
| def __init__(self, models_dir=None, device=None): | |
| # Default to the models/ dir next to this file so standalone .predict() works from any CWD. | |
| md = Path(models_dir) if models_dir else Path(__file__).resolve().parent / "models" | |
| with open(md/"meta.json") as f: self.meta = json.load(f) | |
| self.LES = self.meta["classes"]; self.CAN = set(self.meta["malignant"]) | |
| self.device = _pick_device(device) | |
| torch.set_num_threads(1) | |
| # Heads run via ONNX Runtime (no lightgbm at inference -> avoids the torch+lightgbm libomp crash). | |
| import onnxruntime as ort | |
| # Store the runtime scalers as plain arrays. Loading sklearn pickles with a | |
| # different sklearn release emits compatibility warnings and can change | |
| # results; mean/scale arrays are all StandardScaler needs at inference. | |
| def load_scaler(name): | |
| with np.load(md/name) as sc: | |
| return sc["mean"].astype("float32"), sc["scale"].astype("float32") | |
| self.scs = load_scaler("scaler_siglip2.npz"); self.scd = load_scaler("scaler_dermlip.npz") | |
| opts = ort.SessionOptions(); opts.intra_op_num_threads = 1 | |
| self.onx_sig = ort.InferenceSession(str(md/"head_siglip2.onnx"), sess_options=opts, providers=["CPUExecutionProvider"]) | |
| self.onx_der = ort.InferenceSession(str(md/"head_dermlip.onnx"), sess_options=opts, providers=["CPUExecutionProvider"]) | |
| self.cls_sig = self.meta["class_order_siglip2"]; self.cls_der = self.meta["class_order_dermlip"] | |
| opp = md/"operating_point.json" | |
| if opp.exists(): | |
| with open(opp) as f: self.op = json.load(f) | |
| else: | |
| self.op = {"threshold_high_sensitivity":0.08,"uncertain_gate":{"max_prob_below":0.45,"entropy_above":0.98,"margin_below":0.10},"operating_point":{}} | |
| # DermLIP (dermatology CLIP) — load FIRST (open_clip before transformers avoids the co-load segfault) | |
| import open_clip | |
| self.der_model, _, self.der_pre = open_clip.create_model_and_transforms(self.meta["embedders"]["dermlip"]["open_clip"]) | |
| self.der_model = self.der_model.eval().to(self.device) | |
| # SigLIP2 (general VLM) — L2-normalized pooled image features | |
| from transformers import AutoModel, AutoProcessor | |
| sid = self.meta["embedders"]["siglip2"]["hf_id"] | |
| self.sig_proc = AutoProcessor.from_pretrained(sid) | |
| self.sig_model = AutoModel.from_pretrained(sid).eval().to(self.device) | |
| def embed(self, img: Image.Image): | |
| img = ImageOps.exif_transpose(img).convert("RGB") | |
| inp = self.sig_proc(images=img, return_tensors="pt").to(self.device) | |
| out = self.sig_model.get_image_features(**inp) | |
| out = out.pooler_output if hasattr(out, "pooler_output") else out # SigLIP2 pooled features | |
| f = out.float().cpu().numpy()[0]; f = f/max(float(np.linalg.norm(f)), 1e-12) # L2 normalize (== bank) | |
| x = self.der_pre(img).unsqueeze(0).to(self.device) | |
| g = self.der_model.encode_image(x).float().cpu().numpy()[0] # raw | |
| return f.astype("float32"), g.astype("float32") | |
| def _onnx_proba(self, sc, sess, cls_order, vec): | |
| mean, scale = sc | |
| Xs = ((vec.reshape(1,-1) - mean) / scale).astype("float32") | |
| outs = sess.run(None, {"input": Xs}) | |
| prob = next(np.asarray(o)[0] for o in outs if np.asarray(o).ndim==2 and np.asarray(o).shape[1]==len(self.LES)) | |
| full = np.zeros(len(self.LES)) | |
| for j,c in enumerate(cls_order): full[self.LES.index(c)] = prob[j] | |
| return full | |
| def predict(self, img: Image.Image, red_flags=None) -> dict: | |
| """Returns the contract-style four-state decision for DermaFusion. `red_flags` is a dict of | |
| safety_override_inputs (e.g. {"changing": true}); any True can only ESCALATE to referral.""" | |
| f,g = self.embed(img) | |
| Psig = self._onnx_proba(self.scs, self.onx_sig, self.cls_sig, f) | |
| Pder = self._onnx_proba(self.scd, self.onx_der, self.cls_der, g) | |
| P = 0.5*(Psig + Pder) | |
| order = np.argsort(P)[::-1] | |
| per_class = [{"code": self.LES[i], "label": CLASS_LABELS[self.LES[i]][0], | |
| "malignancy": CLASS_LABELS[self.LES[i]][1], "prob": round(float(P[i]),4)} for i in order] | |
| hr = float(sum(P[self.LES.index(c)] for c in self.CAN)) | |
| maxp = float(P[order[0]]); margin = float(P[order[0]] - P[order[1]]) | |
| pc = np.clip(P,1e-9,1); entropy = float(-(pc*np.log(pc)).sum()) | |
| entropy_norm = entropy / math.log(len(self.LES)) # normalize to [0,1] so the gate matches operating_point.json | |
| T = float(self.op.get("threshold_high_sensitivity",0.08)) | |
| ug = self.op.get("uncertain_gate",{"max_prob_below":0.45,"entropy_above":0.98,"margin_below":0.10}) | |
| # ---- contract four-state decision (decision is the source of truth, NOT a cosmetic band) ---- | |
| # abstain if: low top-class confidence, near-uniform spread (normalized entropy), or a near-tie top-2 margin. | |
| if maxp < ug.get("max_prob_below",0.45) or entropy_norm > ug.get("entropy_above",0.98) or margin < ug.get("margin_below",0.10): | |
| decision = "uncertain_or_ood" | |
| elif hr >= T: | |
| decision = "high_risk_triage_warning" | |
| else: | |
| decision = "low_high_risk_signal" | |
| # ---- red-flag safety override: can ONLY escalate to referral, never lower ---- | |
| reasons = [k for k,v in (red_flags or {}).items() if v] | |
| escalated = False | |
| if reasons and decision in ("low_high_risk_signal","uncertain_or_ood"): | |
| decision = "high_risk_triage_warning"; escalated = True | |
| ACTION = { | |
| "high_risk_triage_warning": "See a dermatologist for evaluation.", | |
| "low_high_risk_signal": "No strong high-risk signal in this photo — this is NOT an all-clear. Keep watching this spot and see a doctor if it changes.", | |
| "uncertain_or_ood": "This image can't be assessed confidently. If the spot is new or changing, see a doctor regardless.", | |
| "poor_quality_retake": "Couldn't read this photo well enough — retake in better light.", | |
| } | |
| return { | |
| "model": self.meta["model"], "version": self.meta["version"], | |
| "decision": decision, "recommended_action": ACTION[decision], | |
| "high_risk_score": round(hr,4), "threshold": T, | |
| "top_likelihood": per_class[0], # most likely class as a LIKELIHOOD, never a diagnosis | |
| "per_class": per_class, | |
| "max_prob": round(maxp,4), "confidence_margin": round(margin,4), "entropy": round(entropy,4), | |
| "safety_override_reasons": reasons, "escalated_by_red_flag": escalated, | |
| "operating_point": self.op.get("operating_point",{}), | |
| "disclaimer": self.meta["disclaimer"], | |
| } | |
| _ENGINE=None | |
| def get_engine(models_dir=None, device=None): | |
| global _ENGINE | |
| if _ENGINE is None: _ENGINE=DermaFusionEngine(models_dir, device) | |
| return _ENGINE | |