app.py

"""
Medical Image Analysis Tool
Providers:
  - Groq         → meta-llama/llama-4-scout-17b-16e-instruct  (needs GROQ_API_KEY)
  - HF Router    → zai-org/GLM-4.6V-Flash                     (needs HF_TOKEN)

Requirements:
    pip install gradio groq openai pillow ddgs

API Keys:
    - Groq (FREE): https://console.groq.com/keys
    - HF Token:    https://huggingface.co/settings/tokens
"""

import os
import base64
from io import BytesIO

from PIL import Image as PILImage
from ddgs import DDGS
from openai import OpenAI
import gradio as gr


# ---------------------------------------------------------------------------
# Model Registry
# ---------------------------------------------------------------------------

MODEL_OPTIONS = {
    "[Groq]  Llama-4-Scout-17B":           {
        "provider":   "groq",
        "model_id":   "meta-llama/llama-4-scout-17b-16e-instruct",
        "base_url":   "https://api.groq.com/openai/v1",
        "key_env":    "GROQ_API_KEY",
    },
    "[HF]    GLM-4.6V-Flash  · novita":    {
        "provider":   "hf",
        "model_id":   "zai-org/GLM-4.6V-Flash:novita",
        "base_url":   "https://router.huggingface.co/v1",
        "key_env":    "HF_TOKEN",
    },
    "[HF]    GLM-4.6V-Flash  · zai-org":   {
        "provider":   "hf",
        "model_id":   "zai-org/GLM-4.6V-Flash:zai-org",
        "base_url":   "https://router.huggingface.co/v1",
        "key_env":    "HF_TOKEN",
    },
    "[HF]    GLM-4.6V-Flash  · fastest":   {
        "provider":   "hf",
        "model_id":   "zai-org/GLM-4.6V-Flash:fastest",
        "base_url":   "https://router.huggingface.co/v1",
        "key_env":    "HF_TOKEN",
    },
}

DEFAULT_MODEL = "[Groq]  Llama-4-Scout-17B"

ANALYSIS_PROMPT = """
You are a highly skilled medical imaging expert with extensive knowledge in radiology and diagnostic imaging and pcos detetcion. Analyze the medical image and structure your response as follows:

### 1. Image Type & Region
- Identify imaging modality (X-ray/MRI/CT/Ultrasound/etc.).
- Specify anatomical region and positioning.
- Evaluate image quality and technical adequacy.

### 2. Key Findings
- Highlight primary observations systematically.
- Identify potential abnormalities with detailed descriptions.
- Include measurements and densities where relevant.

### 3. Diagnostic Assessment
- Provide primary diagnosis with confidence level.
- List differential diagnoses ranked by likelihood.
- Support each diagnosis with observed evidence.
- Highlight critical/urgent findings.

### 4. Patient-Friendly Explanation
- Simplify findings in clear, non-technical language.
- Avoid medical jargon or provide easy definitions.
- Include relatable visual analogies.

Note: Research references will be appended separately.
Ensure a structured and medically accurate response using clear markdown formatting.
"""


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def pil_to_base64(img: PILImage.Image) -> str:
    buf = BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def search_medical_literature(query: str, max_results: int = 3) -> str:
    try:
        with DDGS() as ddgs:
            results = list(ddgs.text(query, max_results=max_results))
        if not results:
            return "_No search results found._"
        lines = []
        for i, r in enumerate(results, 1):
            lines.append(f"**[{i}] [{r['title']}]({r['href']})**")
            lines.append(f"{r['body']}\n")
        return "\n".join(lines)
    except Exception as e:
        return f"_Search unavailable: {e}_"


def extract_diagnosis_keywords(text: str) -> str:
    lines = text.splitlines()
    in_section, keywords = False, []
    for line in lines:
        if "Diagnostic Assessment" in line:
            in_section = True
            continue
        if in_section:
            if line.startswith("###"):
                break
            stripped = line.strip("- *").strip()
            if stripped:
                keywords.append(stripped)
        if len(keywords) >= 2:
            break
    if keywords:
        return f"medical imaging {' '.join(keywords[:2])} diagnosis treatment protocol"
    return "radiology diagnostic imaging findings treatment guidelines"


# ---------------------------------------------------------------------------
# Core Analysis
# ---------------------------------------------------------------------------

def analyze_medical_image(
    image: PILImage.Image,
    model_choice: str,
    groq_key: str,
    hf_token: str,
) -> str:
    if image is None:
        return "⚠️ Please upload a medical image."

    cfg = MODEL_OPTIONS[model_choice]
    provider = cfg["provider"]

    # Resolve API key: UI input → env var fallback
    api_key = (groq_key.strip() if provider == "groq" else hf_token.strip()) \
              or os.environ.get(cfg["key_env"], "")

    if not api_key:
        key_url = "https://console.groq.com/keys" if provider == "groq" \
                  else "https://huggingface.co/settings/tokens"
        label = "Groq API key" if provider == "groq" else "HuggingFace token"
        return (
            f"⚠️ **{label} required for this model.**\n\n"
            f"Get yours (free) at: {key_url}"
        )

    # Both providers use the same OpenAI-compatible interface
    client = OpenAI(api_key=api_key, base_url=cfg["base_url"])

    img_b64 = pil_to_base64(image)

    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{img_b64}"},
                },
                {"type": "text", "text": ANALYSIS_PROMPT},
            ],
        }
    ]

    try:
        resp = client.chat.completions.create(
            model=cfg["model_id"],
            messages=messages,
            max_tokens=2048,
        )
        analysis = resp.choices[0].message.content

    except Exception as e:
        err = str(e)
        tip = ""
        if "401" in err or "unauthorized" in err.lower():
            tip = "\n\n**Fix:** API key is invalid or expired."
        elif "429" in err or "rate" in err.lower():
            tip = "\n\n**Fix:** Rate limit hit — wait a moment and retry."
        elif "404" in err:
            tip = "\n\n**Fix:** Model endpoint not found. Try a different variant."
        elif "decommissioned" in err.lower():
            tip = "\n\n**Fix:** Model was deprecated — switch to another option."
        return f"⚠️ Inference error:\n```\n{err}\n```{tip}"

    # DuckDuckGo research pass
    search_query = extract_diagnosis_keywords(analysis)
    research = search_medical_literature(search_query)

    return f"""
{analysis}

---

### 5. Research Context
> *Search query: `{search_query}`*

{research}

---
*⚠️ Disclaimer: For informational purposes only. Always consult a qualified healthcare professional.*
""".strip()


# ---------------------------------------------------------------------------
# UI helpers — show/hide key fields based on selected model
# ---------------------------------------------------------------------------

def update_key_visibility(model_choice: str):
    provider = MODEL_OPTIONS[model_choice]["provider"]
    return (
        gr.update(visible=(provider == "groq")),
        gr.update(visible=(provider == "hf")),
    )


# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

with gr.Blocks(title="Medical Image Analysis", theme=gr.themes.Soft()) as demo:

    gr.Markdown("""
    # 🩺 Medical Image Analysis Tool
    Upload a medical image (X-ray, MRI, CT, Ultrasound) for AI-powered analysis.
    Choose your model below — both providers offer a **free tier**.
    """)

    with gr.Row():

        # ── Left column ──────────────────────────────────────────────────
        with gr.Column(scale=1):

            model_selector = gr.Dropdown(
                choices=list(MODEL_OPTIONS.keys()),
                value=DEFAULT_MODEL,
                label="🤖 Model",
            )

            groq_key_input = gr.Textbox(
                label="🔑 Groq API Key",
                placeholder="gsk_xxxxxxxxxxxxxxxxxxxxxxxx",
                type="password",
                value=os.environ.get("GROQ_API_KEY", ""),
                info="Free at console.groq.com/keys",
                visible=True,
            )

            hf_token_input = gr.Textbox(
                label="🔑 HuggingFace Token",
                placeholder="hf_xxxxxxxxxxxxxxxxxxxxxxxx",
                type="password",
                value=os.environ.get("HF_TOKEN", ""),
                info="Free at huggingface.co/settings/tokens",
                visible=False,
            )

            input_image = gr.Image(type="pil", label="📤 Upload Medical Image")

            analyze_button = gr.Button("🔍 Analyze Image", variant="primary", size="lg")

            gr.Markdown("""
            **Model notes:**
            | Model | Provider | Key needed |
            |---|---|---|
            | Llama-4-Scout-17B | Groq | `GROQ_API_KEY` |
            | GLM-4.6V-Flash · novita | HF Router | `HF_TOKEN` |
            | GLM-4.6V-Flash · zai-org | HF Router | `HF_TOKEN` |
            | GLM-4.6V-Flash · fastest | HF Router | `HF_TOKEN` |

            GLM-4.6V-Flash variants differ by inference backend — try **novita** first.
            """)

        # ── Right column ─────────────────────────────────────────────────
        with gr.Column(scale=2):
            output_report = gr.Markdown(
                value="_Your analysis report will appear here._",
                label="📋 Analysis Report",
            )

    # Auto-show/hide the correct key field when model changes
    model_selector.change(
        fn=update_key_visibility,
        inputs=model_selector,
        outputs=[groq_key_input, hf_token_input],
    )

    analyze_button.click(
        fn=analyze_medical_image,
        inputs=[input_image, model_selector, groq_key_input, hf_token_input],
        outputs=output_report,
    )

if __name__ == "__main__":
    demo.launch(debug=True, share=True)