Update app.py

app.py

@@ -1,322 +1,214 @@
-# app.py
-"""
-Hugging Face Space / Gradio app for Acne Type/Severity Classification + Chatbot (Mistral)
-- Input: Image URL (user provides)
-- Model: loads a Hugging Face image-classification model (default recommended checkpoint)
-- Explanation: returns textual explanation for predicted acne label
-- Chatbot: uses Mistral Chat Completions API (user supplies API key)
-"""
-
 import os
 import io
 import requests
 from PIL import Image
+from huggingface_hub import InferenceApi
 import gradio as gr
 
-# Transformers imports
-from transformers import pipeline, AutoImageProcessor, AutoModelForImageClassification
-import torch
-
-# --------------------------
-# CONFIG: choose model here
-# --------------------------
-# You can swap this to any HF image-classification checkpoint that supports acne/skin labels.
-MODEL_NAME = "imfarzanansari/skintelligent-acne"  # recommended default (acne severity)
-# Fallbacks (used if primary model fails to load):
-FALLBACK_MODELS = [
-    "naamalia23/acne-severity-classification",
-    "Tanishq77/skin-condition-classifier"
-]
-
-# Mistral API end-point (chat completions)
-MISTRAL_CHAT_URL = "https://api.mistral.ai/v1/chat/completions"
-
-# --------------------------
-# Utility helpers
-# --------------------------
-def load_model(model_name):
-    """
-    Try to load HF image-classification pipeline for model_name.
-    Returns a pipeline object or raises.
-    """
-    try:
-        device = 0 if torch.cuda.is_available() else -1
-        classifier = pipeline("image-classification", model=model_name, device=device)
-        return classifier
-    except Exception as e:
-        raise RuntimeError(f"Failed to load model {model_name}: {e}")
-
-# Try to load the chosen model, fallback if necessary
-classifier = None
-loaded_model_name = None
-load_errors = []
-try:
-    classifier = load_model(MODEL_NAME)
-    loaded_model_name = MODEL_NAME
-except Exception as e:
-    load_errors.append(str(e))
-    for alt in FALLBACK_MODELS:
-        try:
-            classifier = load_model(alt)
-            loaded_model_name = alt
-            break
-        except Exception as e2:
-            load_errors.append(str(e2))
-
-if classifier is None:
-    # If no model loaded, app will still start but classification will return helpful error
-    print("WARNING: No classification model loaded. Errors:", load_errors)
-
-# --------------------------
-# Simple textual explanations for common labels
-# (Customize / extend as needed for your model's label set)
-# --------------------------
-EXPLANATION_BANK = {
-    # examples for acne severity labels (modify as per the model labels)
-    "Level -1: Clear Skin": "No active acne detected. Skin appears clear. Maintain gentle cleansing and sunscreen.",
-    "Level 0: Occasional Spots": "Occasional pimples or spots. Often manageable with over-the-counter topical treatments (benzoyl peroxide, salicylic acid).",
-    "Level 1: Mild Acne": "Mild acne with comedones (whiteheads/blackheads) and a few papules. Use topical retinoids, gentle cleanser; seek dermatologist if persistent.",
-    "Level 2: Moderate Acne": "Moderate acne with inflammatory papules and pustules. Prescription topical or oral treatments may be needed. See dermatologist for tailored therapy.",
-    "Level 3: Severe Acne": "Severe inflammatory acne, possibly nodules or cysts. Early dermatologist consultation is strongly recommended; systemic therapy may be needed.",
-    "Level 4: Very Severe Acne": "Very severe acne with widespread nodules/cysts or scarring. Urgent dermatologist evaluation required for systemic and procedural options.",
-    # fallback generic labels
-    "acne": "Signs of acne detected. Severity and subtype should be confirmed by a clinician. Usual treatments range from topical care to systemic medications depending on severity.",
-    "mild": "Mild acne. Start with gentle skincare and OTC active ingredients; consult dermatologist if it doesn't improve.",
-    "moderate": "Moderate acne. Dermatology visit recommended; topical and/or oral therapies may be indicated.",
-    "severe": "Severe acne. Dermatologist assessment needed; potential for scarring and systemic therapy."
-}
-
-def get_explanation_for_label(label):
-    # direct match
-    if label in EXPLANATION_BANK:
-        return EXPLANATION_BANK[label]
-    # case-insensitive partial match
-    ll = label.lower()
-    for k, v in EXPLANATION_BANK.items():
-        if k.lower() in ll or ll in k.lower():
-            return v
-    # fallback
-    return ("Detected label: {}. This model's label indicates acne or a related skin condition. "
-            "If you want a more specific explanation, fine-tune the EXPLANATION_BANK for your model's labels.").format(label)
+# Configuration (expect these to be set as environment variables in the Space)
+HF_MODEL_ID = os.environ.get("HF_MODEL_ID", "imfarzanansari/skintelligent-acne")
+HF_API_TOKEN = os.environ.get("HF_API_TOKEN")  # optional (better rate-limits if provided)
+MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY")  # set this in Secrets for the Space
+MISTRAL_MODEL = os.environ.get("MISTRAL_MODEL", "mistral-large-latest")
+
+# Initialize Hugging Face Inference API client (image classification model)
+inference = InferenceApi(repo_id=HF_MODEL_ID, token=HF_API_TOKEN)
 
-# --------------------------
-# Image download and prepare
-# --------------------------
-def load_image_from_url(url):
+
+def fetch_image_from_url(url: str) -> Image.Image:
+    resp = requests.get(url, timeout=15)
+    resp.raise_for_status()
+    return Image.open(io.BytesIO(resp.content)).convert("RGB")
+
+
+def classify_acne(image_url: str):
     try:
-        resp = requests.get(url, timeout=10)
-        resp.raise_for_status()
-        img = Image.open(io.BytesIO(resp.content)).convert("RGB")
-        return img
+        img = fetch_image_from_url(image_url)
     except Exception as e:
-        raise RuntimeError(f"Failed to fetch image from URL: {e}")
-
-# --------------------------
-# Classification function (used by Gradio)
-# --------------------------
-def classify_image_from_url(image_url):
-    if classifier is None:
-        return {
-            "status": "error",
-            "message": "No model available. Check server logs or swap MODEL_NAME to a valid checkpoint."
-        }
-
-    # fetch image
+        return {"error": f"Failed to fetch image: {e}"}
+
+    # Convert to bytes for inference
+    buf = io.BytesIO()
+    img.save(buf, format="JPEG")
+    image_bytes = buf.getvalue()
+
+    # Call the HF Inference API model
     try:
-        img = load_image_from_url(image_url)
+        hf_response = inference(inputs=image_bytes)
     except Exception as e:
-        return {"status": "error", "message": str(e)}
+        return {"error": f"Model inference failed: {e}"}
+
+    # hf_response is usually a list of dicts with 'label' and 'score' keys
+    top_preds = []
+    if isinstance(hf_response, dict) and "error" in hf_response:
+        return {"error": hf_response.get("error")}
 
-    # run inference (pipeline returns list of dicts)
     try:
-        preds = classifier(img, top_k=3)
-    except Exception as e:
-        return {"status": "error", "message": f"Model inference failed: {e}"}
-
-    # normalize output format
-    # preds -> list like [{"label": "Level 1: Mild", "score": 0.91}, ...]
-    top = preds[0]
-    label = top.get("label", str(top))
-    score = float(top.get("score", 0.0))
-
-    explanation = get_explanation_for_label(label)
-
-    # construct a concise structured response for the UI
-    response = {
-        "status": "ok",
-        "model": loaded_model_name or "none",
-        "label": label,
-        "score": round(score, 4),
+        for item in hf_response[:5]:
+            label = item.get("label")
+            score = item.get("score")
+            top_preds.append({"label": label, "score": float(score)})
+    except Exception:
+        # fallback: cast the response to string
+        top_preds = [{"label": str(hf_response), "score": 1.0}]
+
+    # Build a short explanation using the LLM (Mistral)
+    explanation = generate_explanation_with_mistral(top_preds)
+
+    return {
+        "image": img,
+        "predictions": top_preds,
         "explanation": explanation,
-        "top_predictions": preds
     }
-    return response
 
-# --------------------------
-# Mistral Chatbot integration
-# --------------------------
-def call_mistral_chat(api_key: str, messages: list, model: str = "mistral-small-latest", stream: bool = False):
-    """
-    Call the Mistral Chat Completions endpoint.
-    messages: a list of dicts, e.g. [{"role":"user", "content":"..."}]
-    returns response text (single string) or raise.
+
+def generate_explanation_with_mistral(preds):
+    """Call Mistral Chat Completions API to produce a human-readable explanation.
+    preds: list of {label, score}
     """
-    if not api_key:
-        raise RuntimeError("Mistral API key is required for chatbot.")
-    headers = {
-        "Authorization": f"Bearer {api_key}",
-        "Content-Type": "application/json"
-    }
+    if not MISTRAL_API_KEY:
+        # Fallback explanation if no API key is set
+        lines = [f"{p['label']}: {p['score']*100:.1f}%" for p in preds]
+        return (
+            "\n".join(lines)
+            + "\n\n(Note: No Mistral API key provided — enable MISTRAL_API_KEY to get a richer explanation.)"
+        )
+
+    system_prompt = (
+        "You are a helpful dermatology-aware assistant. Provide a concise, non-diagnostic explanation "
+        "of the acne findings listed. For each predicted label give: (1) what it means, (2) likely causes/triggers, "
+        "(3) simple self-care suggestions, and (4) when to see a dermatologist. Keep language suitable for laypeople. "
+        "Always include a clear disclaimer that this is not medical advice."
+    )
+
+    user_message = "Detected acne predictions:\n"
+    for p in preds:
+        user_message += f"- {p['label']}: {p['score']*100:.1f}%\n"
+
+    user_message += (
+        "\nPlease explain the above concisely and practically. Use bullet points for clarity and finish with a short disclaimer."
+    )
+
     body = {
-        "model": model,
-        "messages": messages
+        "model": MISTRAL_MODEL,
+        "messages": [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_message},
+        ],
+        "max_tokens": 400,
+        "temperature": 0.2,
+    }
+
+    headers = {
+        "Authorization": f"Bearer {MISTRAL_API_KEY}",
+        "Content-Type": "application/json",
     }
+
     try:
-        r = requests.post(MISTRAL_CHAT_URL, json=body, headers=headers, timeout=30)
+        r = requests.post("https://api.mistral.ai/v1/chat/completions", json=body, headers=headers, timeout=30)
         r.raise_for_status()
-        data = r.json()
-        # parse returned content
-        choices = data.get("choices", [])
-        if len(choices) > 0:
-            # some Mistral endpoints put the message under choices[0]["message"]["content"]
-            msg = choices[0].get("message", {}) or {}
-            content = msg.get("content") or choices[0].get("text") or ""
-            # content may be a string or dict; ensure string
-            if isinstance(content, dict):
-                # join parts if necessary
-                content = content.get("text", str(content))
-            return content
-        # fallback flat text
-        return data.get("text", str(data))
+        j = r.json()
+        # Mistral returns choices; extract the assistant reply
+        text = ""
+        # nested structure may vary; try common fields
+        if isinstance(j, dict):
+            # try several common patterns
+            if "choices" in j and len(j["choices"]) > 0:
+                text = j["choices"][0].get("message", {}).get("content", "")
+            elif "output" in j and isinstance(j["output"], str):
+                text = j["output"]
+            elif "text" in j:
+                text = j["text"]
+        if not text:
+            text = str(j)
+        return text
     except Exception as e:
-        raise RuntimeError(f"Mistral API call failed: {e}")
+        return f"(LLM explanation unavailable: {e})"
 
-# --------------------------
-# Gradio UI callbacks
-# --------------------------
-# Keep simple conversation state via closure
-chat_history = []
 
-def classify_and_prepare_context(image_url):
-    """
-    Runs classification and returns structured outputs plus
-    a "context" text that the chatbot can use (label + explanation).
-    """
-    result = classify_image_from_url(image_url)
-    if result.get("status") != "ok":
-        return None, result.get("message", "Unknown error")
-    # Build context summary
-    context_summary = (
-        f"Detected acne label: {result['label']} (confidence {result['score']}). "
-        f"Explanation: {result['explanation']}"
-    )
-    return result, context_summary
+# --- Gradio UI ---
 
-def chat_with_context(mistral_api_key, user_message, context_summary, model_name="mistral-small-latest"):
-    """
-    Send conversation to Mistral with context prepended.
-    Returns assistant reply (string).
-    """
-    if not mistral_api_key:
-        return "Please provide your Mistral API key (in the Mistral API Key box) to use the chatbot."
-
-    # maintain in-memory chat history for nicer flow
-    # We will prepend a system message + context on every call to give the model grounding
-    system_msg = {
-        "role": "system",
-        "content": (
-            "You are a helpful, concise dermatology assistant. Use clinical but accessible language. "
-            "Base your answers on standard dermatology practice. If you are unsure, recommend seeing a dermatologist."
-        )
+def run_analysis(image_url: str):
+    result = classify_acne(image_url)
+    if "error" in result:
+        return result["error"], None, "", None
+
+    # Image for display
+    img = result["image"]
+
+    # Build human-friendly prediction lines
+    pred_lines = "\n".join([f"{p['label']} — {p['score']*100:.1f}%" for p in result["predictions"]])
+
+    explanation = result["explanation"]
+
+    return None, img, pred_lines, explanation
+
+
+# Chatbot: simple Mistral-backed chat using the same API key
+
+def mistral_chat(messages):
+    """messages is a list of (role, content) tuples where role is 'user' or 'assistant'."""
+    if not MISTRAL_API_KEY:
+        return "Mistral API key not set. Set MISTRAL_API_KEY environment variable in your Space."
+
+    body = {
+        "model": MISTRAL_MODEL,
+        "messages": [{"role": r, "content": c} for r, c in messages],
+        "max_tokens": 400,
+        "temperature": 0.2,
     }
-    context_msg = {"role": "system", "content": context_summary}
-    user_msg = {"role": "user", "content": user_message}
+    headers = {"Authorization": f"Bearer {MISTRAL_API_KEY}", "Content-Type": "application/json"}
+    r = requests.post("https://api.mistral.ai/v1/chat/completions", json=body, headers=headers, timeout=30)
+    r.raise_for_status()
+    j = r.json()
+    # extract assistant message
+    if "choices" in j and len(j["choices"]) > 0:
+        return j["choices"][0].get("message", {}).get("content", "")
+    # fallback
+    return str(j)
 
-    messages = [system_msg, context_msg, user_msg]
 
-    try:
-        reply = call_mistral_chat(api_key=mistral_api_key, messages=messages, model=model_name)
-    except Exception as e:
-        return f"[Chat error] {e}"
-    return reply
-
-# --------------------------
-# Build Gradio app layout
-# --------------------------
-with gr.Blocks(theme=gr.themes.Default(), title="Acne Classifier + Mistral Chatbot") as demo:
-    gr.Markdown("## Acne Type/Severity Classifier + Chatbot\n"
-                "Paste an **image URL** (a photo of the face/skin area). The app will classify acne type/severity "
-                "and provide an explanation. Use the chatbot (Mistral) to ask follow-up questions about the diagnosis, "
-                "treatments, and next steps. **You must provide your Mistral API key** to use the chatbot.")
+with gr.Blocks(title="Acne Type Classifier & Advisor") as demo:
+    gr.Markdown("# Acne Type Classifier — Upload image URL, get classification + explanation")
     with gr.Row():
-        with gr.Column(scale=2):
-            image_url_input = gr.Textbox(label="Image URL", placeholder="https://...", lines=1)
-            load_and_classify_btn = gr.Button("Load & Classify")
-            image_output = gr.Image(label="Loaded Image", type="pil")
-            model_info = gr.Textbox(value=f"Model loaded: {loaded_model_name or 'None'}", label="Model info", interactive=False)
-            results_box = gr.JSON(label="Classification Result (structured)", interactive=False)
         with gr.Column(scale=1):
-            mistral_key_input = gr.Textbox(label="Mistral API Key", placeholder="sk-...", type="password")
-            gr.Markdown("### Chatbot about detected acne")
-            chat_output = gr.Chatbot(label="Dermatology Assistant")
-            user_msg_input = gr.Textbox(placeholder="Ask about the detected acne...", label="Your question")
-            send_btn = gr.Button("Send")
-
-    # classify button action
-    def on_classify_click(url):
-        if not url or url.strip() == "":
-            return None, {"status":"error","message":"Please paste an image URL"}, None
-        # show image
-        try:
-            img = load_image_from_url(url)
-        except Exception as e:
-            return None, {"status":"error","message":str(e)}, None
-        result, context = classify_and_prepare_context(url)
-        if result is None:
-            return img, {"status":"error","message": context}, None
-        # Preload the chat history reset
-        global chat_history
-        chat_history = []
-        # Return image to display, JSON results, and put context into a hidden area via gr.State if needed
-        return img, result, context
-
-    load_and_classify_btn.click(on_classify_click, inputs=[image_url_input], outputs=[image_output, results_box, gr.State()])
-
-    # chat send action
-    def on_send_click(mkey, user_text, last_context):
-        if not last_context:
-            return gr.update(), "Please classify an image first (use the Load & Classify button)."
-        if not user_text or user_text.strip() == "":
-            return gr.update(), "Please type a question."
-        # call Mistral
-        reply = chat_with_context(mkey, user_text, last_context)
-        # Append to chat_history and return
-        global chat_history
-        chat_history.append(("User", user_text))
-        chat_history.append(("Assistant", reply))
-        # Format chat_history as list of tuples for gr.Chatbot
-        formatted = [(u, a) for u, a in zip(chat_history[::2], chat_history[1::2])]
-        return formatted, ""
-    # Note: gr.State will hold the latest context_summary; as a simple approach, we pass last output results_box['explanation'] as context.
-    # But Gradio's .click binding above returned a third value (context) which is not stored here; for simplicity we re-run classification to extract context.
-    # We'll implement a small wrapper to grab the context from the results_box JSON client-side.
-    # For clarity and reliability in Spaces, recommend wiring a hidden State; here we accept the user to paste Mistral key and ask after classifying.
-
-    send_btn.click(
-        fn=lambda key, text, context_summary: (
-            # return updated chat and cleared input
-            chat_with_context(key, text, context_summary),
-            ""
-        ),
-        inputs=[mistral_key_input, user_msg_input, results_box],
-        outputs=[chat_output, user_msg_input]
-    )
+            url_in = gr.Textbox(label="Image URL", placeholder="https://...jpg")
+            analyze_btn = gr.Button("Analyze Image")
+            error_out = gr.Textbox(label="Errors", interactive=False)
+
+        with gr.Column(scale=1):
+            image_out = gr.Image(type="pil", label="Image")
+            preds_out = gr.Textbox(label="Top Predictions", interactive=False)
+            explanation_out = gr.Textbox(label="Explanation (from LLM)", interactive=False)
+
+    analyze_btn.click(run_analysis, inputs=[url_in], outputs=[error_out, image_out, preds_out, explanation_out])
+
+    gr.Markdown("---\n## Ask follow-up questions (chat)")
+
+    chatbot = gr.Chatbot()
+    user_msg = gr.Textbox(label="Your question")
+    send_btn = gr.Button("Send")
+
+    # maintain history in a state
+    history = gr.State([])
+
+    def on_send(user_text, history):
+        history = history or []
+        history.append(("user", user_text))
+        assistant = mistral_chat(history)
+        history.append(("assistant", assistant))
+        # convert to gr.Chatbot format (list of [user, assistant] pairs)
+        chat_display = [[h[1] if h[0]=="user" else "", h[1] if h[0]=="assistant" else ""] for h in history if h[0] in ("user","assistant")]
+        # better format: pairwise
+        pairs = []
+        for i in range(0, len(history), 2):
+            u = history[i][1] if i < len(history) else ""
+            a = history[i+1][1] if i+1 < len(history) else ""
+            pairs.append([u, a])
+        return pairs, history
+
+    send_btn.click(on_send, inputs=[user_msg, history], outputs=[chatbot, history])
 
-    gr.Markdown("**Notes & Tips**:\n\n"
-                "- If pipeline/model loading fails on startup, change `MODEL_NAME` to another HF checkpoint and restart the Space.\n"
-                "- For best results: clear, well-lit closeup photos of acne lesions give higher accuracy.\n"
-                "- This app provides informational assistance only — not a medical diagnosis. Encourage users to consult a dermatologist for medical decisions.")
 
-# Launch
 if __name__ == "__main__":
     demo.launch()