Update app.py

app.py

@@ -1,187 +1,147 @@
-import os
-import io
-import base64
-import requests
-from PIL import Image
 import gradio as gr
+from huggingface_hub import InferenceApi
 from transformers import pipeline
+import requests
+from PIL import Image
+import io
+import os
+import openai
 
-# ----------------------------
-# Configuration
-# ----------------------------
-MODEL_ID = os.environ.get("HF_MODEL_ID", "imfarzanansari/skintelligent-acne")
-MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY")
-MISTRAL_MODEL = os.environ.get("MISTRAL_MODEL", "mistral-large-latest")
+# -----------------------------
+# CONFIGURATION
+# -----------------------------
+# Hugging Face model for acne classification
+MODEL_ID = "imfarzanansari/skintelligent-acne"
 
-# Load local acne classifier pipeline
+# Use local pipeline for image classification (faster and more stable)
 classifier = pipeline("image-classification", model=MODEL_ID)
 
-# ----------------------------
-# Helper functions
-# ----------------------------
-def fetch_image_from_url(url: str) -> Image.Image:
-    """Download and return image from URL."""
-    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):
-    """Run acne classification and explanation."""
-    try:
-        img = fetch_image_from_url(image_url)
-    except Exception as e:
-        return {"error": f"Failed to fetch image: {e}"}
+# Set your Mistral API key (via environment variable)
+openai.api_key = os.getenv("MISTRAL_API_KEY")
 
+# -----------------------------
+# HELPER FUNCTIONS
+# -----------------------------
+def classify_acne(image_url):
     try:
-        preds = classifier(img)
+        response = requests.get(image_url)
+        img = Image.open(io.BytesIO(response.content)).convert("RGB")
     except Exception as e:
-        return {"error": f"Model inference failed: {e}"}
+        return "❌ Could not load image. Please check the URL.", "", None
 
-    # Keep top 5 predictions
-    top_preds = preds[:5]
-    explanation = generate_explanation_with_mistral(top_preds)
+    # Run the acne classification
+    preds = classifier(img)
+    if not preds:
+        return "No prediction.", "", img
 
-    return {"image": img, "predictions": top_preds, "explanation": explanation}
+    top_pred = preds[0]["label"]
+    score = preds[0]["score"]
 
-def generate_explanation_with_mistral(preds):
-    """Generate explanation using Mistral API."""
-    if not MISTRAL_API_KEY:
-        lines = [f"{p['label']}: {p['score']*100:.1f}%" for p in preds]
-        return (
-            "\n".join(lines)
-            + "\n\n(Note: Mistral API key not set — enable MISTRAL_API_KEY to get full explanations.)"
-        )
+    # Explanation text
+    explanation = explain_acne_type(top_pred)
 
-    system_prompt = (
-        "You are a helpful dermatology assistant. Explain each acne type simply: "
-        "1) what it means, 2) likely causes, 3) basic care, 4) when to see a dermatologist. "
-        "End with a disclaimer that this is not medical advice."
-    )
+    result_text = f"**Detected Acne Type:** {top_pred}\n\n**Confidence:** {score:.2f}"
+    return result_text, explanation, img
 
-    user_message = "Detected acne predictions:\n"
-    for p in preds:
-        user_message += f"- {p['label']}: {p['score']*100:.1f}%\n"
 
-    user_message += "\nExplain these findings clearly and briefly."
-
-    headers = {
-        "Authorization": f"Bearer {MISTRAL_API_KEY}",
-        "Content-Type": "application/json",
+def explain_acne_type(acne_type):
+    explanations = {
+        "Blackheads": "Blackheads are open comedones caused by clogged hair follicles with sebum and dead skin. They appear black due to oxidation.",
+        "Whiteheads": "Whiteheads are closed comedones formed when pores are clogged with oil and dead skin but remain closed at the surface.",
+        "Papules": "Papules are small, red, inflamed bumps without visible pus. They often result from irritated or infected clogged pores.",
+        "Pustules": "Pustules are pus-filled pimples with a white or yellow center. They can be tender and are often caused by bacterial infection.",
+        "Nodules": "Nodules are large, painful lumps deep under the skin caused by severe inflammation and infection in clogged pores.",
+        "Cysts": "Cysts are severe acne lesions filled with pus and can cause scarring if not treated properly.",
     }
+    return explanations.get(acne_type, "This acne type is uncommon or not specifically defined in the dataset.")
 
-    body = {
-        "model": MISTRAL_MODEL,
-        "messages": [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_message},
-        ],
-        "max_tokens": 400,
-        "temperature": 0.2,
-    }
 
+def query_acne_info(acne_type, user_query):
+    if not user_query.strip():
+        return "Please enter a question."
     try:
-        r = requests.post(
-            "https://api.mistral.ai/v1/chat/completions",
-            json=body,
-            headers=headers,
-            timeout=30,
+        prompt = f"You are an expert dermatologist. The user has acne type '{acne_type}'. Answer this query:\n{user_query}"
+        completion = openai.ChatCompletion.create(
+            model="mistral-tiny",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.6,
         )
-        r.raise_for_status()
-        j = r.json()
-        if "choices" in j and len(j["choices"]) > 0:
-            return j["choices"][0]["message"]["content"]
-        return str(j)
+        return completion.choices[0].message["content"]
     except Exception as e:
-        return f"(LLM explanation unavailable: {e})"
-
-# ----------------------------
-# Chatbot
-# ----------------------------
-def mistral_chat(messages):
-    """Simple chat interface using Mistral API."""
-    if not MISTRAL_API_KEY:
-        return "Mistral API key not set."
-
-    body = {
-        "model": MISTRAL_MODEL,
-        "messages": [{"role": r, "content": c} for r, c in messages],
-        "max_tokens": 400,
-        "temperature": 0.3,
-    }
-    headers = {
-        "Authorization": f"Bearer {MISTRAL_API_KEY}",
-        "Content-Type": "application/json",
-    }
-
-    try:
-        r = requests.post(
-            "https://api.mistral.ai/v1/chat/completions",
-            json=body,
-            headers=headers,
-            timeout=30,
-        )
-        r.raise_for_status()
-        j = r.json()
-        if "choices" in j and len(j["choices"]) > 0:
-            return j["choices"][0]["message"]["content"]
-        return str(j)
-    except Exception as e:
-        return f"(Chat unavailable: {e})"
-
-# ----------------------------
-# Gradio UI
-# ----------------------------
-def run_analysis(image_url: str):
-    result = classify_acne(image_url)
-    if "error" in result:
-        return result["error"], None, "", None
-
-    img = result["image"]
-    preds_str = "\n".join(
-        [f"{p['label']} — {p['score']*100:.1f}%" for p in result["predictions"]]
+        return f"Error: {str(e)}"
+
+
+# -----------------------------
+# GRADIO INTERFACE
+# -----------------------------
+with gr.Blocks(theme=gr.themes.Soft(), title="Acne Type Classifier & Chatbot") as demo:
+    gr.Markdown(
+        """
+        # 🧴 Acne Type Classifier & Dermatology Assistant  
+        Upload or paste the URL of an acne image to detect its type.  
+        Then ask any query about the detected acne type using the chatbot below.
+        """
     )
-    return None, img, preds_str, result["explanation"]
-
-with gr.Blocks(title="Acne Type Classifier & Chat") as demo:
-    gr.Markdown("# Acne Type Classifier — Enter Image URL and Get Insights")
 
     with gr.Row():
-        with gr.Column(scale=1):
-            url_in = gr.Textbox(label="Image URL", placeholder="https://...jpg")
-            analyze_btn = gr.Button("Analyze Image")
-            error_out = gr.Textbox(label="Error", interactive=False)
+        image_url = gr.Textbox(label="🔗 Enter Image URL", placeholder="Paste image URL here...")
+        submit_btn = gr.Button("Classify", variant="primary")
 
+    with gr.Row():
         with gr.Column(scale=1):
-            image_out = gr.Image(type="pil", label="Input Image")
-            preds_out = gr.Textbox(label="Predictions", interactive=False)
-            explanation_out = gr.Textbox(label="Explanation", interactive=False)
-
-    analyze_btn.click(
-        run_analysis, inputs=[url_in], outputs=[error_out, image_out, preds_out, explanation_out]
+            image_output = gr.Image(label="Uploaded Image", type="pil")
+
+        with gr.Column(scale=2):
+            result_box = gr.Markdown(label="Prediction Result", elem_classes="big-box")
+            explanation_box = gr.Textbox(
+                label="Acne Explanation",
+                lines=6,
+                interactive=False,
+                elem_classes="big-box"
+            )
+
+    # Chatbot section
+    with gr.Accordion("💬 Ask Dermatology Chatbot", open=True):
+        with gr.Row():
+            user_query = gr.Textbox(
+                label="Enter your query about the detected acne",
+                placeholder="e.g., What is the best treatment for cystic acne?",
+                lines=2,
+            )
+        with gr.Row():
+            chat_response = gr.Textbox(
+                label="Chatbot Response",
+                lines=6,
+                interactive=False,
+                elem_classes="big-box"
+            )
+        chat_btn = gr.Button("Ask Chatbot", variant="secondary")
+
+    # Button functionality
+    submit_btn.click(
+        classify_acne,
+        inputs=[image_url],
+        outputs=[result_box, explanation_box, image_output],
     )
 
-    gr.Markdown("---\n## Ask a Question About Acne")
-
-    chatbot = gr.Chatbot()
-    user_msg = gr.Textbox(label="Your Question")
-    send_btn = gr.Button("Send")
-
-    history = gr.State([])
-
-    def on_send(user_text, history):
-        history = history or []
-        history.append(("user", user_text))
-        reply = mistral_chat(history)
-        history.append(("assistant", reply))
+    chat_btn.click(
+        query_acne_info,
+        inputs=[result_box, user_query],
+        outputs=[chat_response],
+    )
 
-        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
+    gr.Markdown(
+        "#### ⚕️ Disclaimer: This app provides general information and should not replace professional medical advice."
+    )
 
-    send_btn.click(on_send, inputs=[user_msg, history], outputs=[chatbot, history])
+# Custom CSS to enlarge boxes
+demo.css = """
+.big-box textarea, .big-box pre, .big-box .wrap {
+    height: 220px !important;
+    font-size: 16px;
+}
+"""
 
+# Launch app
 if __name__ == "__main__":
     demo.launch()