File size: 30,500 Bytes
8253efc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d8c2226
8253efc
 
 
753c901
8253efc
 
 
 
753c901
8253efc
 
 
b842c0f
1c964a9
 
 
b842c0f
8253efc
 
 
 
 
753c901
8253efc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b842c0f
753c901
 
 
8253efc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
753c901
b842c0f
 
 
 
 
8253efc
 
 
753c901
8253efc
 
 
 
 
 
 
 
 
 
b842c0f
8253efc
b842c0f
 
 
 
 
 
 
 
 
 
 
8253efc
 
b842c0f
 
 
 
 
 
8253efc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
753c901
8253efc
 
753c901
 
 
 
 
 
 
 
 
 
1c964a9
 
 
 
753c901
 
 
8253efc
 
753c901
8253efc
 
 
 
 
 
 
 
 
753c901
8253efc
 
 
 
 
 
 
 
753c901
8253efc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
753c901
8253efc
753c901
8253efc
 
753c901
8253efc
 
 
753c901
 
8253efc
 
 
 
 
 
 
b842c0f
 
 
 
753c901
8253efc
 
 
 
 
 
 
753c901
8253efc
753c901
8253efc
753c901
8253efc
 
 
 
 
 
 
 
753c901
8253efc
 
 
 
 
 
 
b842c0f
 
 
8253efc
 
 
 
 
 
 
 
 
 
b842c0f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8253efc
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b842c0f
 
 
 
8253efc
 
 
 
753c901
8253efc
 
b842c0f
 
 
1c964a9
b842c0f
 
674c456
2fd84ab
1c964a9
2fd84ab
 
 
 
 
 
 
 
 
 
1c964a9
674c456
1c964a9
 
 
674c456
2fd84ab
 
 
674c456
1c964a9
 
 
 
 
 
 
 
674c456
1c964a9
 
 
 
 
2fd84ab
1c964a9
674c456
 
 
 
 
 
 
2fd84ab
674c456
2fd84ab
674c456
 
 
2fd84ab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c964a9
 
 
 
 
2fd84ab
 
 
 
 
 
 
 
1c964a9
 
 
 
2fd84ab
1c964a9
 
 
 
 
 
 
 
 
2fd84ab
1c964a9
 
 
2fd84ab
1c964a9
b842c0f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1c964a9
 
 
 
 
 
 
 
 
 
 
 
b842c0f
 
 
 
8253efc
 
 
 
 
 
 
 
 
 
 
 
753c901
 
 
8253efc
 
 
 
753c901
8253efc
753c901
8253efc
 
 
 
 
753c901
8253efc
 
753c901
8253efc
753c901
8253efc
 
 
 
 
b842c0f
 
753c901
8253efc
 
 
753c901
1c964a9
 
b842c0f
1c964a9
 
 
 
753c901
b842c0f
 
1c964a9
 
 
753c901
2fd84ab
 
b842c0f
753c901
b842c0f
 
8253efc
 
753c901
8253efc
 
 
753c901
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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
"""
Skin Disease Detection β€” Demo
------------------------------
A Gradio app for Hugging Face Spaces with two tabs:
  1. Image Diagnosis β€” upload (or pick a test image) and get a classifier
     prediction plus a plain-language, urgency-aware explanation.
  2. Chat Assistant β€” a multi-turn triage chatbot. It can accept an image,
     asks the kind of follow-up questions a real triage intake would, and
     after enough context gives a preliminary assessment with urgency-tiered
     guidance.

IMPORTANT: This remains a research/portfolio project, not a validated
medical device. It has not been clinically evaluated or regulatory
cleared. It should never be the sole basis for a health decision β€” that
framing shows up once, clearly, rather than as a repeated warning block
after every message, but it's still true and still matters.
"""

import os

import gradio as gr
import spaces
import requests
from PIL import Image
from transformers import pipeline

# ---------------------------------------------------------------------------
# Config β€” override any of these via Space "Variables and secrets"
# ---------------------------------------------------------------------------
IMAGE_MODEL_ID = os.environ.get("IMAGE_MODEL_ID", "Anwarkh1/Skin_Cancer-Image_Classification")
# Hugging Face retired the old api-inference.huggingface.co Serverless API in
# favor of the "Inference Providers" router (OpenAI-compatible). Any chat
# model listed at https://huggingface.co/models?inference_provider=hf-inference
# works here.
CHAT_MODEL_ID = os.environ.get("CHAT_MODEL_ID", "openai/gpt-oss-120b")
HF_TOKEN = os.environ.get("HF_TOKEN")  # add as a Space SECRET to enable the chatbot
ROUTER_URL = "https://router.huggingface.co/v1/chat/completions"

PAGE_DISCLAIMER = (
    "This tool is a research/portfolio project, not a medical device. It hasn't "
    "been clinically validated. Use it to get informed and organize your thoughts "
    "before a real appointment β€” not as a substitute for one."
)

# Optional second classifier specifically for pigmented-lesion cases (HAM10000-based,
# where melanoma is its own clean label). Useful because broader taxonomies like
# DermNet lump "Melanoma Skin Cancer, Nevi and Moles" into one label β€” this lets a
# mole/nevus-flagged image get a second, more specific opinion. Leave unset to disable.
LESION_MODEL_ID = os.environ.get("LESION_MODEL_ID", "")

# ---------------------------------------------------------------------------
# Urgency tiering β€” keyword-based rather than an exact-match label dict, so it
# works across different classifier taxonomies (the original 7-class HAM10000
# labels, the 23-class DermNet labels, or any future swap) without needing to
# hardcode every exact label string. Checked in priority order.
# ---------------------------------------------------------------------------
URGENT_KEYWORDS = [
    "melanoma", "malignant", "carcinoma", "actinic keratos",
    "bullous", "lupus", "systemic", "vasculitis", "cellulitis",
]
PROMPT_DOCTOR_KEYWORDS = [
    "fungal", "fungus", "candidiasis", "tinea", "ringworm",
    "scabies", "lyme", "infestation", "wart", "molluscum",
    "viral", "herpes", "hpv", "std", "exanthem", "drug eruption",
    "impetigo", "bacterial",
]
# Anything not matching the above falls through to "general_care" β€” typically
# manageable/chronic conditions (acne, eczema, psoriasis, hives, hair loss,
# benign tumors, common moles, etc.) where general skin-care guidance fits.

SEVERITY = {
    "urgent": {"label": "High", "emoji": "πŸ”΄", "note": "recommend seeing a dermatologist/doctor soon"},
    "prompt_doctor_visit": {"label": "Medium", "emoji": "🟑", "note": "needs a proper diagnosis/prescription, not an emergency"},
    "general_care": {"label": "Low", "emoji": "🟒", "note": "typically manageable with general skin care"},
}


def classify_tier(label):
    """Map a classifier label to an urgency tier via keyword matching."""
    l = label.lower()
    for kw in URGENT_KEYWORDS:
        if kw in l:
            return "urgent"
    for kw in PROMPT_DOCTOR_KEYWORDS:
        if kw in l:
            return "prompt_doctor_visit"
    return "general_care"


def severity_line(label):
    tier = classify_tier(label)
    s = SEVERITY[tier]
    return f"{s['emoji']} **Severity: {s['label']}** β€” {s['note']}"


# Friendly one-line descriptions for labels we recognize exactly (both the
# original 7-class HAM10000 set and common DermNet-style labels). Purely
# cosmetic β€” if a label isn't here, the tier-based guidance still works fine
# without a description.
KNOWN_DESCRIPTIONS = {
    "actinic keratoses": "A rough, scaly patch caused by sun damage.",
    "basal cell carcinoma": "The most common type of skin cancer.",
    "melanoma": "The most serious common type of skin cancer.",
    "benign keratosis-like lesions": "Non-cancerous growths such as seborrheic keratoses or solar lentigines.",
    "dermatofibroma": "A common benign skin nodule, often on the legs.",
    "melanocytic nevi": "Ordinary moles.",
    "vascular lesions": "Blood-vessel-related marks such as angiomas.",
}

SYSTEM_PROMPT = """You are a warm, knowledgeable skin-health triage assistant. You are not a doctor and cannot diagnose anyone with certainty β€” but your job is to be genuinely useful and specific, not to hide behind vague hedging or constant disclaimers.

CONVERSATION STYLE β€” GO DEEPER BEFORE ASSESSING
Have a real back-and-forth, one or two questions at a time. Aim to gather a fuller picture than a quick glance would give β€” typically 5-8 exchanges before your assessment, more if the picture is still unclear, fewer only if the user clearly just wants a fast read or red flags are already obvious. Useful ground to cover (skip anything already answered or clearly irrelevant):
  - Onset: how long they've had it, how it started
  - Evolution: how it's changed over time β€” size, shape, color, texture
  - Symptoms: itching, bleeding, pain, crusting, oozing, discharge, fever, swelling
  - Pattern: is it one spot or spreading, single or multiple, symmetric or not
  - Triggers/exposures: new products, plants, insect bites, contacts, travel, sun exposure
  - History: prior similar episodes, what's been tried already (and whether it helped), relevant personal/family medical history, allergies
  - Location and distribution on the body
The goal of the extra rounds is to genuinely narrow the differential, not to interrogate β€” keep it conversational, and if the user gives a rich answer that covers several of these at once, don't force redundant questions.
If the user shared an image, you'll also receive the image classifier's findings as extra context, including an urgency tier (not shown verbatim to the user) β€” weave that in naturally rather than reading out raw percentages.

GIVING A PRELIMINARY ASSESSMENT
Once you have real depth of context, give a clear, specific preliminary assessment:
- Name the 1-2 most likely categories given everything discussed, in plain language, and briefly say why (which symptoms/pattern point that way). It's fine β€” good, even β€” to be specific about what you think is most likely. What you must NOT do is claim certainty: use "most consistent with" / "this pattern often points to," never "you have X" or "this is definitely X." A real diagnosis needs an in-person exam and sometimes tests (biopsy, culture, bloodwork) that no photo or conversation can replace β€” say this once, naturally, as part of the assessment, not as a repeated disclaimer.
- Always state an explicit severity level as part of the assessment: **Low**, **Medium**, or **High** β€” plus one line on what that means for next steps. Base it on:
  - **High** = cancer-related or ambiguous-but-could-be-cancer findings, autoimmune/systemic/bullous conditions, cellulitis, vasculitis, or any red-flag symptoms (rapid growth, irregular/changing borders, multiple colors, asymmetry, a sore that won't heal, bleeding, spreading redness, fever) β†’ recommend seeing a dermatologist or doctor soon, directly, without softening it into "keep an eye on it." Note explicitly if the classifier's category can't distinguish something dangerous from something benign (e.g. a label that groups melanoma with ordinary moles) β€” when that ambiguity exists, say so and default to High rather than assuming benign.
  - **Medium** = things that need a proper diagnosis and often a prescription to clear up (fungal, viral, bacterial infections, STD-related, infestations like scabies) but aren't emergencies β†’ recommend a doctor visit, explain why self-treatment usually doesn't fully work here. For anything STD-related, stay factual and non-judgmental, and point toward in-person testing rather than guessing from a photo.
  - **Low** = typically manageable conditions (acne, eczema, psoriasis, hives, hair loss, common moles, benign growths, contact dermatitis) with no red flags β†’ explain what it usually is, and give general skin-care guidance: gentle skincare habits, sun protection, not picking/scratching, watching for changes. Still fine to mention a doctor visit is reasonable if they're unsure, worried, or it's not improving.
- Never give specific medications, dosages, or treatment prescriptions β€” general skin-care habits are fine; treating a self- or AI-identified condition with anything beyond general care is not.

URGENT SITUATIONS
If the user describes something urgent β€” rapidly growing lesion, uncontrolled bleeding, signs of spreading infection, fever with a skin issue, severe pain β€” tell them clearly to seek in-person or emergency care promptly, regardless of how many turns you've had.

Keep replies conversational length β€” a few sentences for a question, a bit longer for the assessment itself."""


# ---------------------------------------------------------------------------
# Lazy-loaded models (so the Space boots fast and only loads on first use)
# ---------------------------------------------------------------------------
_classifier = None
_lesion_classifier = None


def get_classifier():
    global _classifier
    if _classifier is None:
        import torch

        device = 0 if torch.cuda.is_available() else -1
        _classifier = pipeline("image-classification", model=IMAGE_MODEL_ID, device=device)
    return _classifier


def get_lesion_classifier():
    """Optional second opinion model for pigmented-lesion cases. Returns None if disabled."""
    global _lesion_classifier
    if not LESION_MODEL_ID:
        return None
    if _lesion_classifier is None:
        import torch

        device = 0 if torch.cuda.is_available() else -1
        _lesion_classifier = pipeline("image-classification", model=LESION_MODEL_ID, device=device)
    return _lesion_classifier


def chat_completion(messages, max_tokens=500):
    """Call the Hugging Face Inference Providers router directly (OpenAI-compatible)."""
    if not HF_TOKEN:
        return None
    resp = requests.post(
        ROUTER_URL,
        headers={"Authorization": f"Bearer {HF_TOKEN}"},
        json={"model": CHAT_MODEL_ID, "messages": messages, "max_tokens": max_tokens},
        timeout=60,
    )
    if not resp.ok:
        # Surface the actual reason from the response body (e.g. "model not
        # supported by any provider") instead of just the generic status line β€”
        # that detail is what actually explains a 400/404, not "Bad Request".
        try:
            detail = resp.json()
        except Exception:
            detail = resp.text
        raise RuntimeError(f"{resp.status_code} error from router for model '{CHAT_MODEL_ID}': {detail}")
    data = resp.json()
    return data["choices"][0]["message"]["content"]


def classify_image(image):
    """Run the primary classifier and return (preds, formatted_lines). If the top
    prediction looks like a mixed/ambiguous pigmented-lesion label and a second
    lesion-specific model is configured, also run that for a second opinion."""
    clf = get_classifier()
    preds = clf(image, top_k=5)
    lines = []
    for p in preds:
        label = p["label"]
        score = p["score"] * 100
        desc = KNOWN_DESCRIPTIONS.get(label.lower(), "")
        desc_txt = f"  \n  _{desc}_" if desc else ""
        lines.append(f"- **{label}** β€” {score:.1f}%{desc_txt}")

    top_label = preds[0]["label"].lower()
    if ("nevi" in top_label or "mole" in top_label) and "melanoma" in top_label:
        lesion_clf = get_lesion_classifier()
        if lesion_clf is not None:
            lesion_preds = lesion_clf(image, top_k=3)
            lines.append("\n_Second opinion from a lesion-specific model (melanoma is a distinct label here):_")
            for p in lesion_preds:
                lines.append(f"- **{p['label']}** β€” {p['score']*100:.1f}%")
            preds = preds + [{"label": f"[lesion-model] {p['label']}", "score": p["score"]} for p in lesion_preds]

    return preds, lines


def build_classifier_context(preds):
    """Turn classifier output into a compact context block for the LLM (not shown raw to the user)."""
    parts = []
    for p in preds[:4]:
        label = p["label"]
        if label.startswith("[lesion-model]"):
            # second-opinion predictions: tier by the underlying label, minus the tag
            bare = label.replace("[lesion-model] ", "")
            tier = classify_tier(bare)
            parts.append(f"{label} ({p['score']*100:.1f}%, second-opinion, {tier})")
            continue
        tier = classify_tier(label)
        desc = KNOWN_DESCRIPTIONS.get(label.lower(), "")
        parts.append(f"{label} ({p['score']*100:.1f}%, {tier}{': ' + desc if desc else ''})")
    return "Image classifier findings for your use (do not read these percentages verbatim): " + "; ".join(parts)


# ---------------------------------------------------------------------------
# Image diagnosis tab (single-shot: upload -> classify -> explain)
# ---------------------------------------------------------------------------
@spaces.GPU
def diagnose_image(image):
    if image is None:
        return "Please upload an image or pick one of the test images first."

    try:
        preds, lines_list = classify_image(image)
    except Exception as e:
        return f"**Model error:** {e}\n\nMake sure `IMAGE_MODEL_ID` is a valid image-classification model."

    lines = ["### πŸ”¬ Classifier prediction\n"] + lines_list
    lines.append(f"\n{severity_line(preds[0]['label'])}")

    if HF_TOKEN:
        context = build_classifier_context(preds)
        prompt = (
            f"{context}\n\n"
            "The user just uploaded a single image with no conversation yet. Give a short "
            "preliminary assessment following your instructions: plain-language explanation of "
            "the top finding(s), an explicit Low/Medium/High severity call, and a brief natural "
            "mention that an in-person exam is what actually confirms things. Note that without "
            "any conversation history you have less context than usual β€” you can still give your "
            "best read, but you may want to suggest the Chat Assistant for a fuller picture if the "
            "case seems ambiguous. Keep it to 5-7 sentences."
        )
        try:
            explanation = chat_completion(
                [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=300,
            )
            lines.append(f"\n### πŸ€– Assessment\n{explanation}")
        except Exception as e:
            lines.append(f"\n_(Assessment unavailable: {e})_")
    else:
        lines.append(
            "\n_Add an `HF_TOKEN` secret to this Space to also get a plain-language "
            "assessment from the chat model here._"
        )

    lines.append(
        "\n---\n_Tip: use the **Chat Assistant** tab for a fuller triage conversation β€” "
        "it can ask follow-up questions and give a more tailored read._"
    )
    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Chat tab β€” multimodal, multi-turn triage
# ---------------------------------------------------------------------------
MIN_ROUNDS_BEFORE_ASSESSMENT = 4  # user turns, not counting the current one


def _extract_text_and_image(message):
    """Normalize gr.ChatInterface(multimodal=True) message input into (text, image_path_or_None)."""
    if isinstance(message, dict):
        text = message.get("text", "") or ""
        files = message.get("files") or []
        image_path = files[0] if files else None
        return text, image_path
    return str(message), None


def _count_user_turns(messages):
    return sum(1 for m in messages if m.get("role") == "user")


def _stage_note(user_turn_count):
    """Injected each turn so multi-round behavior is enforced by code, not just
    requested once in the system prompt β€” the model gets a fresh reminder of
    what stage the conversation is in on every single call."""
    if user_turn_count < MIN_ROUNDS_BEFORE_ASSESSMENT:
        return (
            f"[Internal note β€” not from the user: this is exchange {user_turn_count} of at least "
            f"{MIN_ROUNDS_BEFORE_ASSESSMENT} before a full assessment, UNLESS the user has already "
            "described clear red-flag symptoms (rapid growth, bleeding, severe pain, spreading "
            "infection, fever) in which case address that urgently right now regardless of turn "
            "count. Otherwise, continue asking focused follow-up questions β€” don't jump to a full "
            "Low/Medium/High assessment yet.]"
        )
    return (
        f"[Internal note β€” not from the user: this is exchange {user_turn_count}, enough rounds "
        "have happened. If you have a reasonably clear picture, give your preliminary assessment "
        "now with an explicit Low/Medium/High severity call. If something important is still "
        "genuinely unclear, you may ask one more focused question first, but don't stall further.]"
    )


def chat_respond(message, history):
    if not HF_TOKEN:
        return (
            "Chat isn't configured yet β€” add an `HF_TOKEN` secret to this Space "
            "(Settings β†’ Variables and secrets) to enable the chatbot."
        )

    text, image_path = _extract_text_and_image(message)

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    for turn in history:
        if isinstance(turn, dict):
            role = turn.get("role")
            content = turn.get("content")
            if isinstance(content, str):
                messages.append({"role": role, "content": content})
        else:
            user_msg, bot_msg = turn
            if isinstance(user_msg, str):
                messages.append({"role": "user", "content": user_msg})
            if bot_msg:
                messages.append({"role": "assistant", "content": bot_msg})

    user_content = text
    if image_path:
        try:
            img = Image.open(image_path)
            preds, _ = classify_image(img)
            context = build_classifier_context(preds)
            user_content = f"{context}\n\nUser's message: {text or '(no message, just shared an image)'}"
        except Exception as e:
            user_content = f"(Image analysis failed: {e})\n\nUser's message: {text}"

    # current turn counts too, since the model is about to respond to it
    current_turn_count = _count_user_turns(messages) + 1
    stage_note = _stage_note(current_turn_count)
    messages.append({"role": "user", "content": f"{stage_note}\n\n{user_content}"})

    try:
        return chat_completion(messages, max_tokens=500)
    except Exception as e:
        return f"Sorry, I hit an error talking to the model: {e}"


# ---------------------------------------------------------------------------
# Model Performance tab data loading
# ---------------------------------------------------------------------------
ASSETS_DIR = "assets"




def load_performance_reports():
    """Scan assets/ for model performance data in two modes:
    1. Full reports: a *_metrics.json paired with its confusion-matrix image
       and any other same-prefix analysis images (threshold curves, cluster
       breakdowns, etc.).
    2. Image-only reports: any *_confusion_matrix.{png,jpg,jpeg} that doesn't
       have a matching metrics.json yet -- shown with its charts but no
       numeric table, so you don't have to wait on exporting metrics.json
       just to see the images you already generated.
    Returns a list of dicts; empty list if nothing's been uploaded at all.
    """
    import json
 
    reports = []
    if not os.path.isdir(ASSETS_DIR):
        return reports
    all_files = sorted(os.listdir(ASSETS_DIR))
    claimed = set()  # filenames already used by a full (metrics.json-backed) report
 
    # --- Pass 1: full reports (metrics.json + its images) ---
    for fname in all_files:
        if not fname.endswith("_metrics.json"):
            continue
        prefix = fname[: -len("_metrics.json")]
        try:
            with open(os.path.join(ASSETS_DIR, fname)) as f:
                data = json.load(f)
        except Exception:
            continue
 
        image_path = None
        for ext in (".png", ".jpg", ".jpeg"):
            candidate = os.path.join(ASSETS_DIR, f"{prefix}_confusion_matrix{ext}")
            if os.path.isfile(candidate):
                image_path = candidate
                claimed.add(f"{prefix}_confusion_matrix{ext}")
                break
 
        extra_images = []
        for other in all_files:
            if not other.startswith(prefix + "_") or not other.lower().endswith((".png", ".jpg", ".jpeg")):
                continue
            if image_path and other == os.path.basename(image_path):
                continue
            label = other[len(prefix) + 1:].rsplit(".", 1)[0].replace("_", " ").strip().title()
            extra_images.append({"path": os.path.join(ASSETS_DIR, other), "label": label})
            claimed.add(other)
 
        reports.append({"data": data, "image_path": image_path, "extra_images": extra_images})
 
    # --- Pass 2: image-only reports -- any confusion-matrix image not already
    # claimed by a full report above, grouped with its same-prefix siblings ---
    seen_prefixes = set()
    for fname in all_files:
        if fname in claimed or not fname.lower().endswith((".png", ".jpg", ".jpeg")):
            continue
        marker = "_confusion_matrix"
        idx = fname.lower().find(marker)
        if idx == -1:
            continue  # only auto-surface files that look like a confusion matrix
        prefix = fname[:idx]
        if prefix in seen_prefixes:
            continue
        seen_prefixes.add(prefix)
 
        extra_images = []
        for other in all_files:
            if other == fname or other in claimed or not other.startswith(prefix + "_"):
                continue
            if not other.lower().endswith((".png", ".jpg", ".jpeg")):
                continue
            label = other[len(prefix) + 1:].rsplit(".", 1)[0].replace("_", " ").strip().title()
            extra_images.append({"path": os.path.join(ASSETS_DIR, other), "label": label})
            claimed.add(other)
 
        model_name = prefix.replace("_", " ").strip().title() or "Model"
        reports.append({
            "data": {"model_name": model_name, "image_only": True},
            "image_path": os.path.join(ASSETS_DIR, fname),
            "extra_images": extra_images,
        })
 
    return reports



def format_performance_markdown(report):
    data = report["data"]
    lines = [f"### {data.get('model_name', 'Model')}"]
 
    if data.get("image_only"):
        lines.append(
            "_Charts available β€” metrics summary not uploaded yet. Add a matching "
            "`*_metrics.json` to this prefix for the full accuracy/F1/per-class table._"
        )
        return "\n".join(lines)
 
    acc = data.get("accuracy")
    f1m = data.get("f1_macro")
    if acc is not None:
        lines.append(f"**Accuracy:** {acc*100:.1f}%  Β·  **Macro F1:** {f1m*100:.1f}%" if f1m is not None else f"**Accuracy:** {acc*100:.1f}%")
 
    per_class = data.get("per_class", {})
    if per_class:
        lines.append("\n| Class | Precision | Recall | F1 | Support |")
        lines.append("|---|---|---|---|---|")
        for cls, m in per_class.items():
            lines.append(
                f"| {cls} | {m.get('precision', 0)*100:.1f}% | {m.get('recall', 0)*100:.1f}% "
                f"| {m.get('f1', 0)*100:.1f}% | {int(m.get('support', 0))} |"
            )
 
    notes = data.get("notes")
    if notes:
        lines.append(f"\n_Note: {notes}_")
 
    return "\n".join(lines)

ABOUT_MARKDOWN = f"""
## About this project

A research/portfolio demo exploring how far an image classifier + LLM chat
assistant can go toward useful skin-condition triage β€” while being explicit
about what it isn't: a validated medical device.

### What's inside
- **Quick Image Check** β€” upload a photo, get a classifier prediction plus
  an LLM-generated plain-language assessment with an explicit Low/Medium/High
  severity call.
- **Chat Assistant** β€” a multi-turn conversation that asks about onset,
  evolution, symptoms, triggers, and history before giving a more targeted
  preliminary read than a single photo alone could support.
- **Model Performance** β€” the actual test-set metrics for whichever
  fine-tuned model(s) are currently backing this app, not just marketing
  claims about accuracy.

### Models
- **Image classifier:** configurable via the `IMAGE_MODEL_ID` variable.
  Two purpose-built options are included as fine-tuning notebooks in this
  repo: a 7-class pigmented-lesion model (HAM10000/ISIC, ViT-Base) and a
  23-class common-skin-condition model (DermNet, ConvNeXt-Base).
- **Optional second-opinion lesion model:** `LESION_MODEL_ID`, used
  automatically when the primary model's top prediction is an ambiguous
  mixed melanoma/mole-type label, to get a model that treats melanoma as
  its own distinct class.
- **Chat model:** configurable via `CHAT_MODEL_ID`, served through Hugging
  Face's Inference Providers router.

### Severity levels
- πŸ”΄ **High** β€” cancer-related, ambiguous-but-could-be-serious, systemic/
  autoimmune, or red-flag symptoms present. See a doctor soon.
- 🟑 **Medium** β€” likely needs a proper diagnosis and prescription
  (infections, infestations). See a doctor, not an emergency.
- 🟒 **Low** β€” typically manageable with general skin care. Doctor visit
  optional, especially if unsure or it's not improving.

### Known limitations (read this before trusting the output)
- **Not clinically validated.** Benchmark accuracy on a public dataset is
  not the same as real-world diagnostic reliability.
- **Dataset label limitations.** The DermNet taxonomy combines melanoma
  with benign moles under one label β€” the app defaults that whole category
  to High severity rather than pretending to distinguish them.
- **Skin-tone representation.** Public dermatology datasets skew toward
  lighter skin tones; performance on underrepresented skin tones is
  unverified without a dedicated bias audit.
- **No clinical validation, IRB review, or regulatory clearance** has been
  done. This should inform, not replace, professional care.

{PAGE_DISCLAIMER}
"""


# ---------------------------------------------------------------------------
# UI
# ---------------------------------------------------------------------------
EXAMPLES_DIR = "examples"
example_images = []
if os.path.isdir(EXAMPLES_DIR):
    example_images = [
        os.path.join(EXAMPLES_DIR, f)
        for f in sorted(os.listdir(EXAMPLES_DIR))
        if f.lower().endswith((".jpg", ".jpeg", ".png"))
    ]

with gr.Blocks(title="Skin Disease Detection β€” Demo") as demo:
    gr.Markdown("# 🩺 Skin Disease Detection")
    gr.Markdown(f"_{PAGE_DISCLAIMER}_")

    with gr.Tab("πŸ“· Quick Image Check"):
        with gr.Row():
            with gr.Column():
                img_in = gr.Image(type="pil", label="Upload a skin image")
                if example_images:
                    gr.Examples(examples=example_images, inputs=img_in, label="Or try a test image")
                else:
                    gr.Markdown(
                        "_No bundled test images found. Run `scripts/download_examples.py` "
                        "before deploying, or just upload your own image._"
                    )
                analyze_btn = gr.Button("Analyze image", variant="primary")
            with gr.Column():
                result_md = gr.Markdown()
        analyze_btn.click(diagnose_image, inputs=img_in, outputs=result_md)

    with gr.Tab("πŸ’¬ Chat Assistant"):
        gr.ChatInterface(
            fn=chat_respond,
            multimodal=True,
            description=(
                "Talk through what you're noticing, and optionally attach a photo (πŸ“Ž). "
                "The assistant will ask several follow-up questions before giving a preliminary "
                "read with a Low/Medium/High severity call β€” the more you share, the more "
                "targeted that read can be."
            ),
        )

    with gr.Tab("πŸ“Š Model Performance"):
        performance_reports = load_performance_reports()
        if not performance_reports:
            gr.Markdown(
                "_No model performance data yet._\n\n"
                "Run either fine-tuning notebook in this repo, then download the "
                "`*_metrics.json` and `*_confusion_matrix.png` files it produces and place "
                "them in this Space's `assets/` folder. This tab picks them up automatically "
                "on the next restart."
            )
        else:
            for report in performance_reports:
                gr.Markdown(format_performance_markdown(report))
                if report["image_path"]:
                    gr.Image(value=report["image_path"], label="Confusion matrix (test set)", show_label=True)
                for extra in report.get("extra_images", []):
                    gr.Image(value=extra["path"], label=extra["label"], show_label=True)

    with gr.Tab("ℹ️ About"):
        gr.Markdown(ABOUT_MARKDOWN)

    gr.Markdown(
        "---\nBuilt with πŸ€— Transformers + Gradio. "
        f"Image model: `{IMAGE_MODEL_ID}`  Β·  Chat model: `{CHAT_MODEL_ID}`"
    )

if __name__ == "__main__":
    demo.launch()