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