app.py

import os
import io
import json
from datetime import datetime

import streamlit as st
import numpy as np
import cv2
import pandas as pd
from PIL import Image
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.applications.efficientnet import preprocess_input

# ─────────────────────────────────────────────
# CONFIG
# ─────────────────────────────────────────────
IMG_SIZE      = 224
MODEL_PATH    = "model/model.h5"
CLASSES_PATH  = "model/class_names.json"

FALLBACK_CLASSES = ["acne", "blackheads", "dark_spots", "pores", "wrinkles"]

CLASS_INFO = {
    "acne":       {"color": (220, 70,  70),  "emoji": "🔴", "desc": "Inflammatory skin condition causing pimples and cysts"},
    "blackheads": {"color": (60,  60,  60),  "emoji": "⚫", "desc": "Clogged hair follicles with oxidised sebum"},
    "dark_spots": {"color": (160, 100, 50),  "emoji": "🟤", "desc": "Hyperpigmentation caused by excess melanin production"},
    "pores":      {"color": (80,  140, 220), "emoji": "🔵", "desc": "Enlarged or visible pores on the skin surface"},
    "wrinkles":   {"color": (180, 130, 200), "emoji": "🟣", "desc": "Fine lines and wrinkles — common skin aging indicator"},
}

# ─────────────────────────────────────────────
# PAGE SETUP
# ─────────────────────────────────────────────
st.set_page_config(
    page_title="DermalScan AI",
    page_icon="🔬",
    layout="wide",
    initial_sidebar_state="expanded",
)

st.markdown("""
<style>
    .main-title  { font-size:2.4rem; font-weight:800; color:#1a3c6e; margin-bottom:0; }
    .sub-title   { font-size:1rem; color:#5a7a9e; margin-bottom:1.5rem; }
    .disclaimer  { background:#fff8e1; border-left:4px solid #ffc107;
                   padding:12px 16px; border-radius:4px; font-size:0.82rem; color:#555; }
</style>
""", unsafe_allow_html=True)

# ─────────────────────────────────────────────
# LOAD MODEL & CLASS NAMES
# ─────────────────────────────────────────────
@st.cache_resource
def load_keras_model():
    if not os.path.exists(MODEL_PATH):
        return None, False
    m = load_model(MODEL_PATH)
    return m, True

@st.cache_data
def load_class_names():
    if os.path.exists(CLASSES_PATH):
        with open(CLASSES_PATH) as f:
            idx = json.load(f)
        # idx is {class_name: index} — invert to {index: class_name}
        return {v: k for k, v in idx.items()}
    return {i: c for i, c in enumerate(FALLBACK_CLASSES)}

# ─────────────────────────────────────────────
# INFERENCE
# ─────────────────────────────────────────────
def predict(model, face_rgb: np.ndarray, idx_to_class: dict):
    img   = cv2.resize(face_rgb, (IMG_SIZE, IMG_SIZE))
    arr   = preprocess_input(img.astype(np.float32))
    arr   = np.expand_dims(arr, 0)
    probs = model.predict(arr, verbose=0)[0]
    return probs

# ─────────────────────────────────────────────
# FACE DETECTION
# ─────────────────────────────────────────────
def detect_faces(img_rgb: np.ndarray):
    gray    = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2GRAY)
    cascade = cv2.CascadeClassifier(
        cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
    )
    faces = cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(40, 40))
    return list(faces) if len(faces) > 0 else []

def annotate(img_rgb: np.ndarray, faces, predictions, idx_to_class):
    out = img_rgb.copy()
    for i, (x, y, w, h) in enumerate(faces):
        if i >= len(predictions):
            break
        probs     = predictions[i]
        top_idx   = int(np.argmax(probs))
        top_label = idx_to_class[top_idx]
        color     = CLASS_INFO.get(top_label, {"color": (200, 200, 200)})["color"]
        cv2.rectangle(out, (x, y), (x+w, y+h), color, 2)
        text = f"{top_label.replace('_',' ')}: {probs[top_idx]*100:.1f}%"
        (tw, th), _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2)
        cv2.rectangle(out, (x, y-th-10), (x+tw+8, y), color, -1)
        cv2.putText(out, text, (x+4, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255,255,255), 2)
    return out

# ─────────────────────────────────────────────
# SIDEBAR
# ─────────────────────────────────────────────
with st.sidebar:
    st.markdown("### 🔬 DermalScan AI")
    st.markdown("**Model:** EfficientNetB0 (TensorFlow)")
    st.markdown("**Framework:** TensorFlow / Keras")
    st.markdown("**Classes:**")
    for cls, info in CLASS_INFO.items():
        st.markdown(f"{info['emoji']} `{cls}`")
    st.markdown("---")
    st.markdown(
        "<div class='disclaimer'>⚠️ For educational purposes only. "
        "Not a substitute for professional medical diagnosis.</div>",
        unsafe_allow_html=True,
    )

# ─────────────────────────────────────────────
# MAIN
# ─────────────────────────────────────────────
st.markdown("<div class='main-title'>🔬 DermalScan AI</div>", unsafe_allow_html=True)
st.markdown(
    "<div class='sub-title'>EfficientNet-based facial skin condition classifier — "
    "Acne · Blackheads · Dark Spots · Pores · Wrinkles</div>",
    unsafe_allow_html=True,
)

model, model_loaded = load_keras_model()
idx_to_class        = load_class_names()

if not model_loaded:
    st.error(
        "⚠️ `model/model.h5` not found. "
        "Run `DermalScan_Train.ipynb` on Google Colab, then place `model.h5` "
        "and `class_names.json` in the `model/` folder."
    )
    st.stop()

uploaded = st.file_uploader("Upload a face image", type=["jpg", "jpeg", "png"])
if uploaded is None:
    st.info("👆 Upload a facial image to get started.")
    st.stop()

file_bytes = np.frombuffer(uploaded.read(), dtype=np.uint8)
img_bgr    = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
if img_bgr is None:
    st.error("Could not read image. Please upload a valid JPG or PNG.")
    st.stop()

img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

col1, col2 = st.columns(2)
with col1:
    st.subheader("Original Image")
    st.image(img_rgb, use_container_width=True)

faces = detect_faces(img_rgb)
st.write(f"**{len(faces)} face(s) detected**")

predictions = []
for (x, y, w, h) in faces:
    pad  = int(0.1 * max(w, h))
    x1   = max(0, x - pad);             y1 = max(0, y - pad)
    x2   = min(img_rgb.shape[1], x+w+pad); y2 = min(img_rgb.shape[0], y+h+pad)
    crop = img_rgb[y1:y2, x1:x2]
    predictions.append(predict(model, crop, idx_to_class))

annotated = annotate(img_rgb, faces, predictions, idx_to_class)
with col2:
    st.subheader("Analysis Result")
    st.image(annotated, use_container_width=True)

# ── Per-face results ──────────────────────────
if predictions:
    st.subheader("Detailed Predictions")
    for i, probs in enumerate(predictions):
        top_idx   = int(np.argmax(probs))
        top_label = idx_to_class[top_idx]
        info      = CLASS_INFO.get(top_label, {"emoji": "⚪", "desc": ""})

        with st.expander(
            f"Face {i+1} — {info['emoji']} **{top_label.replace('_',' ').title()}** "
            f"({probs[top_idx]*100:.1f}% confidence)",
            expanded=True,
        ):
            st.markdown(f"*{info['desc']}*")
            st.markdown("**Probability breakdown:**")
            for j in range(len(probs)):
                cls = idx_to_class.get(j, f"class_{j}")
                st.progress(float(probs[j]), text=f"{cls.replace('_',' ').title()}: {probs[j]*100:.1f}%")

    # DataFrame
    rows = []
    for i, probs in enumerate(predictions):
        row = {"face": i + 1}
        for j in range(len(probs)):
            row[idx_to_class.get(j, f"class_{j}")] = round(float(probs[j]) * 100, 2)
        rows.append(row)
    df = pd.DataFrame(rows)
    st.dataframe(df, use_container_width=True)

    # Downloads
    st.subheader("Export")
    c1, c2 = st.columns(2)
    with c1:
        buf = io.BytesIO()
        Image.fromarray(annotated).save(buf, format="PNG")
        st.download_button("📥 Download Annotated Image", buf.getvalue(),
                           "dermalscan_result.png", "image/png")
    with c2:
        csv = df.to_csv(index=False).encode()
        st.download_button("📊 Download CSV Report", csv,
                           "dermalscan_predictions.csv", "text/csv")
else:
    st.warning("No faces detected. Try a clearer, front-facing photo.")