Update app.py

app.py

@@ -1,47 +1,322 @@
+# 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
 import gradio as gr
-headers = {"Authorization": f"Bearer {MISTRAL_API_KEY}", "Content-Type": "application/json"}
-payload = {
-"model": "mistral-small",
-"messages": [{"role": "system", "content": "You are a dermatologist assistant chatbot."}] +
-[{"role": "user", "content": message}],
-}
-response = requests.post(API_URL, headers=headers, json=payload)
-data = response.json()
-reply = data["choices"][0]["message"]["content"]
-return reply
+
+# 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:
-return f"Error communicating with chatbot: {str(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)
+
+# --------------------------
+# Image download and prepare
+# --------------------------
+def load_image_from_url(url):
+    try:
+        resp = requests.get(url, timeout=10)
+        resp.raise_for_status()
+        img = Image.open(io.BytesIO(resp.content)).convert("RGB")
+        return img
+    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
+    try:
+        img = load_image_from_url(image_url)
+    except Exception as e:
+        return {"status": "error", "message": str(e)}
+
+    # 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)
 
-# Gradio Interface
-def main_app():
-with gr.Blocks() as demo:
-gr.Markdown("# 🧴 AI-Powered Acne Type Classifier")
-gr.Markdown("Enter an image URL of your face or acne region to detect the acne type.")
+    # construct a concise structured response for the UI
+    response = {
+        "status": "ok",
+        "model": loaded_model_name or "none",
+        "label": label,
+        "score": round(score, 4),
+        "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.
+    """
+    if not api_key:
+        raise RuntimeError("Mistral API key is required for chatbot.")
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json"
+    }
+    body = {
+        "model": model,
+        "messages": messages
+    }
+    try:
+        r = requests.post(MISTRAL_CHAT_URL, 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))
+    except Exception as e:
+        raise RuntimeError(f"Mistral API call failed: {e}")
 
-with gr.Row():
-image_url = gr.Textbox(label="Image URL")
-classify_btn = gr.Button("Classify Acne Type")
+# --------------------------
+# 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
 
-image_display = gr.Image(label="Input Image")
-output_text = gr.Markdown()
+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."
+        )
+    }
+    context_msg = {"role": "system", "content": context_summary}
+    user_msg = {"role": "user", "content": user_message}
 
-classify_btn.click(classify_acne, inputs=image_url, outputs=[image_display, output_text])
+    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
 
-gr.Markdown("---")
-gr.Markdown("### 💬 Chat with Dermatologist Assistant")
+# --------------------------
+# 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.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
 
-chatbot = gr.ChatInterface(fn=chat_with_bot, title="Acne Query Chatbot")
+    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.
 
-return demo
+    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]
+    )
 
+    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__":
-app = main_app()
-app.launch()
+    demo.launch()