import json, os import numpy as np from PIL import Image import torch import torch.nn as nn import timm from timm.data import resolve_model_data_config, create_transform from transformers import AutoTokenizer, AutoModel import gradio as gr import ast from huggingface_hub import hf_hub_download SPACE_REPO = os.getenv("SPACE_REPO_NAME", "muruga778/api_for_model") # change if your space id differs def safe_torch_load(filename: str): """ 1) try local file 2) if corrupted -> force-download from Hub cache and load again """ try: print(f"🔎 Loading weights: {filename} (local)") return torch.load(filename, map_location="cpu") except Exception as e: print(f"⚠️ Local load failed for {filename}: {repr(e)}") print("⬇️ Force-downloading from Hugging Face Hub cache...") cached = hf_hub_download( repo_id=SPACE_REPO, repo_type="space", filename=filename, force_download=True, ) print("✅ Downloaded to:", cached, "size(MB)=", os.path.getsize(cached)/1024/1024) return torch.load(cached, map_location="cpu") DEVICE = "cuda" if torch.cuda.is_available() else "cpu" def load_json(path): with open(path, "r") as f: return json.load(f) def clean_state_dict(sd): for key in ["state_dict", "model", "model_state_dict"]: if isinstance(sd, dict) and key in sd and isinstance(sd[key], dict): sd = sd[key] if isinstance(sd, dict) and any(k.startswith("module.") for k in sd.keys()): sd = {k.replace("module.", "", 1): v for k, v in sd.items()} return sd def softmax_np(x): x = x - np.max(x) e = np.exp(x) return e / (np.sum(e) + 1e-9) # --- Triage rules (simple + demo friendly) SEVERITY_BY_LABEL = { "acne": 1, "tinea": 2, "tinea versicolor": 1, "eczema": 2, "urticaria": 2, "psoriasis": 2, "folliculitis": 2, "impetigo": 3, "herpes zoster": 3, "drug rash": 4, "scabies": 3, "unknown": 2 } RED_FLAG_WORDS = [ "fever","breathing","shortness of breath","face","eye","mouth","genital", "severe pain","blister","purple","swelling","rapid","spreading","bleeding" ] def triage(label, conf, text): label_l = (label or "").lower().strip() text_l = (text or "").lower() score = SEVERITY_BY_LABEL.get(label_l, 2) hits = sum(1 for w in RED_FLAG_WORDS if w in text_l) if hits >= 2: score += 2 elif hits == 1: score += 1 if conf < 0.50: score += 1 if conf < 0.35: score += 1 score = int(max(1, min(5, score))) stage = "SELF-CARE / MONITOR" if score <= 2 else ("DOCTOR (24–48h)" if score <= 4 else "URGENT NOW") note = "Not medical advice. If rapidly worsening / fever / face-eye involvement / breathing trouble → seek urgent care." return stage, score, note # ---- Load config + label map CFG = load_json("fusion_config.json") LABEL_MAP = load_json("label_map.json") # Your label_map.json looks like: {"classes":[...], "label2idx":{...}} if isinstance(LABEL_MAP, dict) and "classes" in LABEL_MAP and isinstance(LABEL_MAP["classes"], list): CLASSES = [str(x) for x in LABEL_MAP["classes"]] label2idx = LABEL_MAP.get("label2idx", {c: i for i, c in enumerate(CLASSES)}) # Older possible formats: elif isinstance(LABEL_MAP, dict) and all(isinstance(k, str) and k.isdigit() for k in LABEL_MAP.keys()): # {"0":"eczema", ...} idx2label = {int(k): str(v) for k, v in LABEL_MAP.items()} CLASSES = [idx2label[i] for i in sorted(idx2label.keys())] label2idx = {c: i for i, c in enumerate(CLASSES)} else: # {"eczema": 0, ...} label2idx = {str(k): int(v) for k, v in LABEL_MAP.items()} CLASSES = [c for c, _ in sorted(label2idx.items(), key=lambda x: x[1])] NUM_CLASSES = len(CLASSES) print("✅ NUM_CLASSES:", NUM_CLASSES) print("✅ First labels:", CLASSES[:5]) IMG_BACKBONE = CFG.get("img_backbone", "tf_efficientnetv2_s") IMG_SIZE = int(CFG.get("img_size", 384)) TEXT_MODEL_NAME = CFG.get("text_model_name", "microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext") MAX_LEN = int(CFG.get("max_len", 128)) # ---- Image model img_model = timm.create_model(IMG_BACKBONE, pretrained=False, num_classes=NUM_CLASSES) sd_img = clean_state_dict(safe_torch_load("best_scin_image.pt")) img_model.load_state_dict(sd_img, strict=True) img_model.to(DEVICE).eval() data_cfg = resolve_model_data_config(img_model) data_cfg["input_size"] = (3, IMG_SIZE, IMG_SIZE) img_tfm = create_transform(**data_cfg, is_training=False) # ---- Text model class TextClassifier(nn.Module): def __init__(self, model_name, num_classes, dropout=0.2): super().__init__() self.backbone = AutoModel.from_pretrained(model_name) self.drop = nn.Dropout(dropout) self.head = nn.Linear(self.backbone.config.hidden_size, num_classes) def forward(self, input_ids, attention_mask): out = self.backbone(input_ids=input_ids, attention_mask=attention_mask) feat = out.pooler_output if hasattr(out, "pooler_output") and out.pooler_output is not None else out.last_hidden_state[:, 0] return self.head(self.drop(feat)) tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME) text_model = TextClassifier(TEXT_MODEL_NAME, NUM_CLASSES) sd_txt = clean_state_dict(safe_torch_load("best_scin_text.pt")) text_model.load_state_dict(sd_txt, strict=False) text_model.to(DEVICE).eval() W_IMG = float(CFG.get("fusion_weights", {}).get("image", 0.6)) W_TXT = float(CFG.get("fusion_weights", {}).get("text", 0.4)) s = W_IMG + W_TXT W_IMG, W_TXT = W_IMG / s, W_TXT / s @torch.inference_mode() def predict(image, symptom_text, topk=3): if image is None: return "Upload an image.", "" pil = image.convert("RGB") if hasattr(image, "convert") else Image.open(image).convert("RGB") x_img = img_tfm(pil).unsqueeze(0).to(DEVICE) tok = tokenizer(symptom_text or "", truncation=True, padding="max_length", max_length=MAX_LEN, return_tensors="pt") tok = {k: v.to(DEVICE) for k, v in tok.items()} img_logits = img_model(x_img)[0].detach().float().cpu().numpy() txt_logits = text_model(tok["input_ids"], tok["attention_mask"])[0].detach().float().cpu().numpy() p_img = softmax_np(img_logits) p_txt = softmax_np(txt_logits) p = W_IMG * p_img + W_TXT * p_txt pred_idx = int(np.argmax(p)) pred_label = CLASSES[pred_idx] conf = float(p[pred_idx]) k = min(int(topk), len(CLASSES)) top_idx = np.argsort(-p)[:k] top_lines = [f"{i+1}) {CLASSES[int(ix)]} — {float(p[int(ix)]):.2f}" for i, ix in enumerate(top_idx)] stage, sev_score, note = triage(pred_label, conf, symptom_text) out1 = f"**Prediction:** {pred_label}\n\n**Confidence:** {conf:.2f}\n\n**Triage:** {stage} (score {sev_score}/5)\n\n{note}" out2 = "\n".join(top_lines) return out1, out2 demo = gr.Interface( fn=predict, inputs=[ gr.Image(type="pil", label="Skin image"), gr.Textbox(lines=3, label="Symptoms (text)"), gr.Slider(1, 5, value=3, step=1, label="Top-K"), ], outputs=[ gr.Markdown(label="Result"), gr.Textbox(label="Top-K"), ], title="SmartSkin — SCIN Multimodal (Image + Text Fusion)", description="Demo only. Not medical advice.", ) if __name__ == "__main__": demo.launch()