Spaces:
Sleeping
Sleeping
File size: 3,212 Bytes
4e817fd 8de7c59 a9bb419 4e817fd a9bb419 4e817fd e103386 4e817fd | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | 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 |