Spaces:
Sleeping
Sleeping
| import torch | |
| import gradio as gr | |
| import numpy as np | |
| from PIL import Image | |
| from open_clip import create_model_from_pretrained, get_tokenizer | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| MODEL_NAME = "microsoft/BiomedCLIP-PubMedBERT_256-vit_base_patch16_224" | |
| # ---------------- Load BiomedCLIP ---------------- | |
| print("Loading BiomedCLIP...") | |
| clip_model, preprocess = create_model_from_pretrained(f"hf-hub:{MODEL_NAME}") | |
| tokenizer = get_tokenizer(f"hf-hub:{MODEL_NAME}") | |
| clip_model = clip_model.to(DEVICE).eval() | |
| # ---------------- Load Checkpoint ---------------- | |
| print("Loading Fusion Classifier checkpoint...") | |
| checkpoint = torch.load("fusion_classifier.pth", map_location=DEVICE) | |
| # Some checkpoints store {"state_dict": ...} | |
| if isinstance(checkpoint, dict) and "state_dict" in checkpoint: | |
| state = checkpoint["state_dict"] | |
| else: | |
| state = checkpoint | |
| # ---------------- Determine Architecture From Checkpoint ---------------- | |
| # Required keys must exist | |
| layer1 = state["net.0.weight"].shape # [out, in] e.g. [1024, 1024] | |
| layer2 = state["net.3.weight"].shape # [out, in] e.g. [512, 1024] | |
| layer3 = state["net.6.weight"].shape # [out, in] e.g. [C, 512] | |
| INPUT_DIM = layer1[1] | |
| H1 = layer1[0] | |
| H2 = layer2[0] | |
| NUM_CLASSES = layer3[0] | |
| print(f"Recovered model structure:") | |
| print(f"Input: {INPUT_DIM} -> {H1} -> {H2} -> {NUM_CLASSES}") | |
| # ---------------- Load Class Names ---------------- | |
| try: | |
| classes = checkpoint["classes"] | |
| except: | |
| try: | |
| classes = checkpoint["label_classes"] | |
| except: | |
| print("No class names in checkpoint, loading label_classes.npy") | |
| try: | |
| classes = np.load("label_classes.npy", allow_pickle=True) | |
| except: | |
| print("No label_classes.npy found. Using index labels.") | |
| classes = [f"class_{i}" for i in range(NUM_CLASSES)] | |
| classes = list(classes) | |
| # Safety: if class names count mismatches, fix length | |
| if len(classes) != NUM_CLASSES: | |
| print("Class count mismatch. Adjusting names.") | |
| classes = classes[:NUM_CLASSES] if len(classes) > NUM_CLASSES else classes + [ | |
| f"class_{i}" for i in range(len(classes), NUM_CLASSES) | |
| ] | |
| # ---------------- Build EXACT SAME MODEL ---------------- | |
| class FusionClassifier(torch.nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| self.net = torch.nn.Sequential( | |
| torch.nn.Linear(INPUT_DIM, H1), | |
| torch.nn.ReLU(), | |
| torch.nn.Dropout(0.3), | |
| torch.nn.Linear(H1, H2), | |
| torch.nn.ReLU(), | |
| torch.nn.Dropout(0.3), | |
| torch.nn.Linear(H2, NUM_CLASSES) | |
| ) | |
| def forward(self, x): | |
| return self.net(x) | |
| classifier = FusionClassifier().to(DEVICE) | |
| classifier.load_state_dict(state) | |
| classifier.eval() | |
| # ---------------- Prediction ---------------- | |
| def predict(image, symptoms): | |
| if image is None: | |
| return "Please upload an image" | |
| if not symptoms or symptoms.strip() == "": | |
| symptoms = "no symptom information provided" | |
| img = preprocess(image).unsqueeze(0).to(DEVICE) | |
| text_tok = tokenizer([symptoms]).to(DEVICE) | |
| with torch.no_grad(): | |
| img_feat = clip_model.encode_image(img) | |
| txt_feat = clip_model.encode_text(text_tok) | |
| img_feat /= img_feat.norm(dim=-1, keepdim=True) | |
| txt_feat /= txt_feat.norm(dim=-1, keepdim=True) | |
| fusion = torch.cat([img_feat, txt_feat], dim=1) | |
| logits = classifier(fusion) | |
| probs = torch.softmax(logits, dim=1).cpu().numpy()[0] | |
| idx = probs.argmax() | |
| conf = float(probs[idx]) | |
| label = str(classes[idx]) if idx < len(classes) else "Unknown" | |
| return f"Prediction: {label} | Confidence: {conf:.2f}" | |
| # ---------------- UI ---------------- | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Image(type="pil", label="Upload Skin Image"), | |
| gr.Textbox(label="Describe symptoms (itching, duration, pain, etc.)") | |
| ], | |
| outputs="text", | |
| title="BioLens — Multimodal Skin Condition AI", | |
| description="Upload image + symptoms. Not medical advice." | |
| ) | |
| iface.launch() | |