Spaces:
Sleeping
Sleeping
| import io | |
| import sys | |
| from pathlib import Path | |
| from typing import List, Tuple | |
| from PIL import Image | |
| import importlib.util as _ilu, os as _os | |
| _s = _ilu.spec_from_file_location( | |
| "disease_info", | |
| _os.path.abspath(_os.path.join(_os.path.dirname(__file__), "..", "data", "disease_info.py")) | |
| ) | |
| _m = _ilu.module_from_spec(_s) | |
| _s.loader.exec_module(_m) | |
| DISEASE_INFO = _m.DISEASE_INFO | |
| # --------------------------------------------------------------------------- | |
| # fastai import — handled carefully so the app doesn't crash if GPU is absent | |
| # --------------------------------------------------------------------------- | |
| try: | |
| from fastai.vision.all import load_learner, PILImage | |
| FASTAI_AVAILABLE = True | |
| except ImportError: | |
| FASTAI_AVAILABLE = False | |
| # Path to the .pkl file (project root) | |
| MODEL_PATH = Path(__file__).resolve().parents[2] / "skin_disease_classifier.pkl" | |
| # Target image size the model was trained on | |
| IMAGE_SIZE = 448 | |
| _learner = None # module-level singleton | |
| def load_model(model_path: str = "skin_disease_classifier.pkl") -> bool: | |
| """Load the fastai learner once at startup. Returns True on success.""" | |
| global _learner | |
| if not FASTAI_AVAILABLE: | |
| print("ERROR: fastai is not installed.", file=sys.stderr) | |
| return False | |
| if not MODEL_PATH.exists(): | |
| print(f"ERROR: Model file not found at {MODEL_PATH}", file=sys.stderr) | |
| return False | |
| try: | |
| import pathlib | |
| import platform | |
| print(f"DEBUG: Platform: {platform.system()}", flush=True) | |
| # Patch WindowsPath to PosixPath for Linux loading | |
| pathlib.WindowsPath = pathlib.PosixPath | |
| _learner = load_learner(MODEL_PATH, cpu=True) | |
| print(f"Model loaded successfully from {MODEL_PATH}", flush=True) | |
| return True | |
| except Exception as exc: | |
| import traceback | |
| print(f"ERROR loading model: {exc}", flush=True) | |
| traceback.print_exc() | |
| return False | |
| def is_model_loaded() -> bool: | |
| return _learner is not None | |
| def predict_image(image_bytes: bytes): | |
| import torch | |
| import torchvision.transforms as T | |
| if _learner is None: | |
| raise RuntimeError("Model is not loaded.") | |
| # Build tensor directly — no fastai DataLoader, no threading issues | |
| img = Image.open(io.BytesIO(image_bytes)).convert("RGB") | |
| img = img.resize((IMAGE_SIZE, IMAGE_SIZE)) | |
| tfms = T.Compose([ | |
| T.ToTensor(), | |
| T.Normalize(mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225]), | |
| ]) | |
| tensor = tfms(img).unsqueeze(0) # shape: (1, 3, 448, 448) | |
| _learner.model.eval() | |
| with torch.no_grad(): | |
| logits = _learner.model(tensor) | |
| probs = torch.softmax(logits, dim=1)[0] | |
| vocab = _learner.dls.vocab | |
| label_prob_pairs = list(zip(vocab, probs.tolist())) | |
| label_prob_pairs.sort(key=lambda x: x[1], reverse=True) | |
| top3 = label_prob_pairs[:3] | |
| results = [] | |
| for lbl, prob in top3: | |
| info = DISEASE_INFO.get(lbl, {}) | |
| friendly = info.get("friendly_name", lbl) | |
| results.append({ | |
| "medical_name": lbl, | |
| "friendly_name": friendly, | |
| "confidence": round(float(prob), 4) | |
| }) | |
| return results |