Gradio.py

import torch
import gradio as gr
import torch.nn as nn
import torch.nn.functional as F

from PIL import Image
import torchvision.transforms as transforms
from torchvision.models import resnet50


# ============================================================
# CONFIG
# ============================================================
MODEL_PATH = r"C:\Users\LOQ\Desktop\Oral Diseases Image Classification\checkpoints\best_model.pth"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


# ============================================================
# LOAD MODEL
# ============================================================
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)

CLASS_NAMES = checkpoint["class_names"]
TEST_F1 = checkpoint["test_f1"]

model = resnet50(weights=None)
model.fc = nn.Sequential(
    nn.Dropout(0.3),
    nn.Linear(model.fc.in_features, len(CLASS_NAMES))
)
model.load_state_dict(checkpoint["state_dict"])
model.to(DEVICE)
model.eval()

eval_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])


# ============================================================
# INFERENCE
# ============================================================
def predict(image):
    if image is None:
        return (
            gr.update(value="—", visible=True),
            "—",
            {},
            gr.update(visible=False),
        )

    image = image.convert("RGB")
    tensor = eval_transform(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():
        output = model(tensor)
        probs = F.softmax(output, dim=1)[0]

    index = torch.argmax(probs).item()
    prediction = CLASS_NAMES[index]
    confidence = probs[index].item() * 100

    results = {CLASS_NAMES[i]: float(probs[i]) for i in range(len(CLASS_NAMES))}

    # Build a small verdict badge depending on confidence level
    if confidence >= 85:
        badge = f'<div class="verdict verdict-high">✓ High Confidence — {confidence:.1f}%</div>'
    elif confidence >= 60:
        badge = f'<div class="verdict verdict-mid">! Moderate Confidence — {confidence:.1f}%</div>'
    else:
        badge = f'<div class="verdict verdict-low">? Low Confidence — {confidence:.1f}%</div>'

    return (
        prediction,
        f"{confidence:.2f}%",
        results,
        gr.update(value=badge, visible=True),
    )


def clear_all():
    return None, "—", "—", {}, gr.update(visible=False)


# ============================================================
# STYLING
# ============================================================
CUSTOM_CSS = """

@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&family=JetBrains+Mono:wght@500&display=swap');



:root {

    --bg-primary: #0a0e1a;

    --bg-secondary: #10162a;

    --bg-card: #131a30;

    --border-subtle: #232b45;

    --accent: #14b8a6;

    --accent-soft: #14b8a622;

    --accent-2: #6366f1;

    --text-primary: #e8ecf5;

    --text-secondary: #8892b0;

    --text-muted: #5b6485;

    --radius: 14px;

}



* { font-family: 'Inter', sans-serif !important; }



.gradio-container {

    background: radial-gradient(circle at 10% 0%, #101a33 0%, #080b14 55%, #05070d 100%) !important;

    max-width: 1180px !important;

    margin: 0 auto !important;

}



footer { display: none !important; }



/* ---------- HEADER ---------- */

.app-header {

    padding: 30px 8px 22px 8px;

    border-bottom: 1px solid var(--border-subtle);

    margin-bottom: 26px;

    display: flex;

    align-items: center;

    justify-content: space-between;

}

.app-header .brand {

    display: flex;

    align-items: center;

    gap: 14px;

}

.app-header .logo-badge {

    width: 46px;

    height: 46px;

    border-radius: 12px;

    background: linear-gradient(135deg, var(--accent), var(--accent-2));

    display: flex;

    align-items: center;

    justify-content: center;

    font-size: 22px;

    box-shadow: 0 8px 24px -6px #14b8a655;

    flex-shrink: 0;

}

.app-header h1 {

    font-size: 20px;

    font-weight: 700;

    color: var(--text-primary);

    margin: 0;

    letter-spacing: -0.02em;

}

.app-header p {

    font-size: 13px;

    color: var(--text-muted);

    margin: 2px 0 0 0;

}

.app-header .tag {

    font-size: 11px;

    font-weight: 600;

    color: var(--accent);

    background: var(--accent-soft);

    border: 1px solid #14b8a640;

    padding: 6px 14px;

    border-radius: 999px;

    letter-spacing: 0.03em;

    text-transform: uppercase;

}



/* ---------- CARDS ---------- */

.card {

    background: var(--bg-card) !important;

    border: 1px solid var(--border-subtle) !important;

    border-radius: var(--radius) !important;

    padding: 18px !important;

}

.card-title {

    font-size: 13px;

    font-weight: 600;

    color: var(--text-secondary);

    text-transform: uppercase;

    letter-spacing: 0.04em;

    margin-bottom: 12px;

    display: flex;

    align-items: center;

    gap: 8px;

}

.card-title::before {

    content: "";

    width: 4px;

    height: 14px;

    background: var(--accent);

    border-radius: 2px;

    display: inline-block;

}



/* ---------- UPLOAD ZONE ---------- */

.upload-zone, .upload-zone > div {

    background: var(--bg-card) !important;

    border: 1.5px dashed #2b3454 !important;

    border-radius: var(--radius) !important;

}

.upload-zone:hover {

    border-color: var(--accent) !important;

}



/* ---------- BUTTONS ---------- */

#analyze-btn {

    background: linear-gradient(135deg, #14b8a6, #0d9488) !important;

    color: #05170f !important;

    font-weight: 700 !important;

    border: none !important;

    border-radius: 10px !important;

    box-shadow: 0 10px 24px -8px #14b8a670 !important;

    letter-spacing: 0.01em;

    transition: transform .15s ease, box-shadow .15s ease;

}

#analyze-btn:hover {

    transform: translateY(-1px);

    box-shadow: 0 14px 28px -8px #14b8a690 !important;

}

#clear-btn {

    background: transparent !important;

    color: var(--text-secondary) !important;

    border: 1px solid var(--border-subtle) !important;

    border-radius: 10px !important;

}

#clear-btn:hover {

    border-color: #3a4468 !important;

    color: var(--text-primary) !important;

}



/* ---------- RESULT FIELDS ---------- */

#pred-box textarea, #conf-box textarea {

    background: #0d1326 !important;

    border: 1px solid var(--border-subtle) !important;

    color: var(--text-primary) !important;

    font-weight: 700 !important;

    font-size: 17px !important;

    border-radius: 10px !important;

}

#conf-box textarea {

    color: var(--accent) !important;

    font-family: 'JetBrains Mono', monospace !important;

}

label span {

    color: var(--text-muted) !important;

    font-size: 11.5px !important;

    text-transform: uppercase;

    letter-spacing: 0.05em;

    font-weight: 600 !important;

}



/* ---------- VERDICT BADGE ---------- */

.verdict {

    padding: 10px 16px;

    border-radius: 10px;

    font-size: 13px;

    font-weight: 600;

    text-align: center;

    margin-bottom: 14px;

    border: 1px solid transparent;

}

.verdict-high { background: #14b8a61a; color: #2dd4bf; border-color: #14b8a640; }

.verdict-mid  { background: #f59e0b1a; color: #fbbf24; border-color: #f59e0b40; }

.verdict-low  { background: #ef44441a; color: #f87171; border-color: #ef444440; }



/* ---------- PROBABILITY BARS (gr.Label) ---------- */

.label-wrap {

    background: transparent !important;

    border: none !important;

}

#prob-label .container {

    background: transparent !important;

}



/* ---------- FOOTER ---------- */

.app-footer {

    margin-top: 30px;

    padding: 18px 4px 10px 4px;

    border-top: 1px solid var(--border-subtle);

    display: flex;

    justify-content: space-between;

    align-items: center;

    flex-wrap: wrap;

    gap: 10px;

}

.app-footer .meta {

    font-size: 12px;

    color: var(--text-muted);

    font-family: 'JetBrains Mono', monospace;

}

.app-footer .meta b { color: var(--text-secondary); }

.app-footer .credit {

    font-size: 12px;

    color: var(--text-muted);

}

.app-footer .credit b { color: var(--text-secondary); }



.disclaimer {

    font-size: 11.5px;

    color: var(--text-muted);

    background: #0d132666;

    border: 1px solid var(--border-subtle);

    border-radius: 10px;

    padding: 10px 14px;

    margin-top: 16px;

    line-height: 1.6;

}

"""


# ============================================================
# UI
# ============================================================
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate"),
    css=CUSTOM_CSS,
    title="Oral Disease Classifier"
) as demo:

    gr.HTML(
        f"""

        <div class="app-header">

            <div class="brand">

                <div class="logo-badge">🦷</div>

                <div>

                    <h1>Oral Disease Classification</h1>

                    <p>Computer-vision assisted screening · ResNet50 backbone</p>

                </div>

            </div>

            <div class="tag">Model F1 · {TEST_F1:.3f}</div>

        </div>

        """
    )

    with gr.Row(equal_height=True):

        with gr.Column(scale=5):
            gr.HTML('<div class="card-title">Input Image</div>')
            image_input = gr.Image(
                type="pil",
                label="",
                show_label=False,
                elem_classes="upload-zone",
                height=340,
            )

            with gr.Row():
                clear_btn = gr.Button("Clear", elem_id="clear-btn")
                button = gr.Button("Analyze Image", elem_id="analyze-btn")

            gr.HTML(
                """

                <div class="disclaimer">

                    ⚠ Decision-support tool only. Predictions are generated by an automated

                    model and are not a substitute for professional clinical diagnosis.

                </div>

                """
            )

        with gr.Column(scale=5):
            gr.HTML('<div class="card-title">Analysis Result</div>')

            verdict_html = gr.HTML(visible=False)

            with gr.Row():
                prediction = gr.Textbox(label="Predicted Class", elem_id="pred-box", interactive=False)
                confidence = gr.Textbox(label="Confidence", elem_id="conf-box", interactive=False)

            gr.HTML('<div class="card-title" style="margin-top:6px;">Class Probability Distribution</div>')
            probabilities = gr.Label(
                label="",
                show_label=False,
                elem_id="prob-label",
                num_top_classes=len(CLASS_NAMES),
            )

    gr.HTML(
        f"""

        <div class="app-footer">

            <div class="meta">ARCH · <b>ResNet50</b> &nbsp;|&nbsp; DEVICE · <b>{DEVICE.upper()}</b> &nbsp;|&nbsp; CLASSES · <b>{len(CLASS_NAMES)}</b></div>

            <div class="credit">Built by <b>Nasr Mohamed</b> — AI Engineer &nbsp;·&nbsp; © 2026</div>

        </div>

        """
    )

    button.click(
        predict,
        inputs=image_input,
        outputs=[prediction, confidence, probabilities, verdict_html],
    )

    clear_btn.click(
        clear_all,
        inputs=None,
        outputs=[image_input, prediction, confidence, probabilities, verdict_html],
    )


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