Update app.py

app.py

@@ -1,34 +1,21 @@
 import os
-import re
+import io
+import json
+import base64
 import gradio as gr
-from transformers import pipeline
-import torch
 from openai import OpenAI
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 from pypdf import PdfReader
 
-
-# =========================
-# 1. Device & local vision model
-# =========================
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# local image->text (fast rough caption)
-vision_pipe = pipeline(
-    "image-to-text",
-    model="nlpconnect/vit-gpt2-image-captioning",
-    device=0 if device == "cuda" else -1
-)
-
 # =========================
-# 2. Yunwu / OpenAI-compatible client
+# 1. 云雾 / OpenAI-compatible API 设置
 # =========================
 
 YUNWU_API_KEY = os.environ.get("YUNWU_API_KEY")
 if not YUNWU_API_KEY:
     raise RuntimeError(
-        "YUNWU_API_KEY not set. Add it in HF Space Settings → Variables and secrets."
+        "环境变量 YUNWU_API_KEY 未设置。"
+        "请在 Hugging Face Space 的 Settings → Variables and secrets 中添加。"
     )
 
 client = OpenAI(
@@ -36,39 +23,64 @@ client = OpenAI(
     base_url="https://yunwu.ai/v1"
 )
 
-# =========================
-# 3. Per-step model routing (edit here)
-# =========================
-MODEL_STEP1 = "gpt-5.1"       # strongest reasoning for paper+figure meaning
-MODEL_STEP2 = "deepseek-chat" # cheap/fast for annotation suggestions
-MODEL_STEP3B = "gpt-4o"       # good fluency for plain-language explanation
+VISION_MODEL = "gpt-4o"        # 视觉强模型（可替换成云雾里可用的视觉模型名）
+TEXT_MODEL = "deepseek-chat"   # 纯文本模型（便宜快）
+
+def call_text_llm(prompt: str, model: str = TEXT_MODEL, temperature=0.2, max_tokens=700):
+    resp = client.chat.completions.create(
+        model=model,
+        messages=[
+            {"role": "system",
+             "content": "You are a careful academic assistant. Write clean, natural English."},
+            {"role": "user", "content": prompt}
+        ],
+        temperature=temperature,
+        max_tokens=max_tokens
+    )
+    return resp.choices[0].message.content.strip()
+
+def pil_to_data_url(img: Image.Image, max_side=1400) -> str:
+    img = img.convert("RGB")
+    w, h = img.size
+    scale = min(1.0, max_side / max(w, h))
+    if scale < 1.0:
+        img = img.resize((int(w * scale), int(h * scale)))
+
+    buf = io.BytesIO()
+    img.save(buf, format="PNG")
+    b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:image/png;base64,{b64}"
 
+def call_vision_llm(prompt: str, image: Image.Image, model: str = VISION_MODEL,
+                    temperature=0.2, max_tokens=700, force_json=False):
+    data_url = pil_to_data_url(image)
+
+    # 尝试让模型直接返回 JSON（如果兼容 response_format）
+    kwargs = {}
+    if force_json:
+        kwargs["response_format"] = {"type": "json_object"}
 
-def call_llm(prompt: str,
-             model: str,
-             temperature: float = 0.2,
-             max_tokens: int = 512) -> str:
     resp = client.chat.completions.create(
         model=model,
         messages=[
+            {"role": "system",
+             "content": "You are a careful academic assistant. Read the figure precisely and write natural English."},
             {
-                "role": "system",
-                "content": (
-                    "You are a helpful academic assistant. "
-                    "Be accurate, concrete, and avoid hallucinating numbers. "
-                    "Do not use markdown unless asked."
-                ),
-            },
-            {"role": "user", "content": prompt},
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": prompt},
+                    {"type": "image_url", "image_url": {"url": data_url}}
+                ]
+            }
         ],
         temperature=temperature,
         max_tokens=max_tokens,
+        **kwargs
     )
     return resp.choices[0].message.content.strip()
 
-
 # =========================
-# 4. Utilities: PDF context + output cleaning
+# 2. PDF 上下文提取
 # =========================
 
 def extract_pdf_snippet(pdf_path: str | None, max_chars: int = 2500) -> str:
@@ -87,177 +99,254 @@ def extract_pdf_snippet(pdf_path: str | None, max_chars: int = 2500) -> str:
     except Exception:
         return ""
 
+# =========================
+# 3. Step2 JSON 解析与渲染
+# =========================
 
-def clean_text(text: str) -> str:
+def safe_json_parse(text: str):
     """
-    Remove markdown-ish / AI-ish formatting and obvious repetition.
+    尝试解析模型输出为 JSON。
+    若模型没严格返回纯 JSON，就从文本中截取最外层 {...} 再 parse。
     """
-    if not text:
-        return text
-
-    # remove bold/italic markers
-    text = text.replace("**", "").replace("__", "")
-
-    # remove leading markdown bullets like "- ", "* ", "• "
-    text = re.sub(r"(?m)^\s*[-•*]\s+", "", text)
-
-    # collapse repeated whitespace
-    text = re.sub(r"\n{3,}", "\n\n", text)
-    text = re.sub(r"[ \t]{2,}", " ", text)
-
-    return text.strip()
-
+    try:
+        return json.loads(text)
+    except Exception:
+        # 粗暴截取第一个 { 到最后一个 }
+        start = text.find("{")
+        end = text.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            try:
+                return json.loads(text[start:end+1])
+            except Exception:
+                return None
+        return None
+
+def load_font(size=18):
+    # HF 通常有 DejaVuSans
+    for path in [
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
+        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
+    ]:
+        try:
+            return ImageFont.truetype(path, size=size)
+        except Exception:
+            continue
+    return ImageFont.load_default()
+
+def draw_arrow(draw, x1, y1, x2, y2, color="red", width=3):
+    # 画线
+    draw.line([(x1, y1), (x2, y2)], fill=color, width=width)
+    # 简单箭头三角
+    import math
+    angle = math.atan2(y2 - y1, x2 - x1)
+    head_len = 12
+    head_angle = math.pi / 7
+    p1 = (x2 - head_len * math.cos(angle - head_angle),
+          y2 - head_len * math.sin(angle - head_angle))
+    p2 = (x2 - head_len * math.cos(angle + head_angle),
+          y2 - head_len * math.sin(angle + head_angle))
+    draw.polygon([ (x2, y2), p1, p2 ], fill=color)
+
+def render_annotations(image: Image.Image, ann_json: dict):
+    """
+    ann_json schema:
+    {
+      "annotations": [
+        {"type":"text","text":"...","x":0.5,"y":0.1,"size":18,"color":"red"},
+        {"type":"box","xy":[x1,y1,x2,y2],"outline":"red","width":3},
+        {"type":"arrow","from":[x1,y1],"to":[x2,y2],"color":"red","width":3}
+      ]
+    }
+    Coords normalized 0-1.
+    """
+    if not ann_json or "annotations" not in ann_json:
+        return image
+
+    img = image.convert("RGB").copy()
+    draw = ImageDraw.Draw(img)
+    W, H = img.size
+
+    for ann in ann_json["annotations"]:
+        a_type = ann.get("type", "").lower()
+
+        if a_type == "text":
+            x = ann.get("x", 0.5) * W
+            y = ann.get("y", 0.5) * H
+            txt = ann.get("text", "")
+            color = ann.get("color", "red")
+            size = int(ann.get("size", 18))
+            font = load_font(size=size)
+            # text stroke for visibility
+            draw.text((x, y), txt, fill=color, font=font, stroke_width=2, stroke_fill="white")
+
+        elif a_type == "box":
+            xy = ann.get("xy", [0.1,0.1,0.3,0.3])
+            x1, y1, x2, y2 = xy
+            x1, y1, x2, y2 = x1*W, y1*H, x2*W, y2*H
+            outline = ann.get("outline", "red")
+            width = int(ann.get("width", 3))
+            draw.rectangle([x1,y1,x2,y2], outline=outline, width=width)
+
+        elif a_type == "arrow":
+            f = ann.get("from", [0.2,0.2])
+            t = ann.get("to", [0.4,0.4])
+            x1, y1 = f[0]*W, f[1]*H
+            x2, y2 = t[0]*W, t[1]*H
+            color = ann.get("color", "red")
+            width = int(ann.get("width", 3))
+            draw_arrow(draw, x1,y1,x2,y2, color=color, width=width)
+
+    return img
 
 # =========================
-# 5. Main workflow
+# 4. 主 workflow
 # =========================
 
 def run_workflow(image, keywords, style, pdf_path):
     if image is None:
-        return "", "", ""
+        return "", {}, None, ""
 
-    # Step 0: rough vision caption
-    vision_raw = vision_pipe(image)[0]["generated_text"]
-
-    # contexts
-    kw_text = keywords.strip() if keywords else ""
+    kw_text = (keywords or "").strip()
     pdf_context = extract_pdf_snippet(pdf_path)
 
-    # style control for Step1 only
-    if style == "formal":
-        style_instruction = (
-            "Write in formal academic English suitable for a Results/Figure explanation."
-        )
-    elif style == "fluency":
-        style_instruction = (
-            "Write in smooth, natural academic English, clear and readable."
-        )
-    else:  # simple
-        style_instruction = (
-            "Write clearly with simpler wording, but still accurate."
-        )
-
-    # -------- Step 1: Explanation of what the figure shows (paper+figure) --------
+    # -------- Step1: 论文式解释 ----------
     step1_prompt = f"""
-You are given a scientific figure and some related paper text.
+You are given a scientific figure from a paper.
 
-Rough vision caption of the image:
-\"\"\"{vision_raw}\"\"\"
+Tasks:
+1) Identify what kind of figure it is (bar chart, line plot, conceptual diagram, etc.).
+2) Describe what the x-axis / panels represent and what the y-axis represents in general terms.
+3) State the main pattern or comparison you observe.
+4) Use the paper context to stay on-topic, but do NOT hallucinate exact numbers or p-values.
 
 Figure keywords:
 \"\"\"{kw_text}\"\"\"
 
-Relevant paper context (may be partial/noisy):
+Paper context (may be incomplete/noisy):
 \"\"\"{pdf_context}\"\"\"
 
-Task:
-- Explain what this figure means in 6–8 sentences.
-- Say what is compared on the x-axis (or panels) and what the y-axis measures.
-- Describe the main trends and the scientific takeaway.
-- If the paper context implies a specific mechanism, mention it briefly.
-- Do NOT invent exact numbers, statistics, or p-values.
-- Do NOT describe unrelated parts of the paper.
-
-{style_instruction}
-Return plain text only.
+Write one compact Results-style paragraph (4–6 sentences). Natural academic English.
 """
-    step1_text = call_llm(step1_prompt, model=MODEL_STEP1, max_tokens=520)
-    step1_text = clean_text(step1_text)
+    try:
+        step1_text = call_vision_llm(step1_prompt, image, max_tokens=520)
+    except Exception:
+        step1_text = call_text_llm(step1_prompt, max_tokens=520)
 
-    # -------- Step 2: Suggestions for annotating the figure (less AI-ish) --------
+    # -------- Step2: 输出 JSON 批注计划 ----------
     step2_prompt = f"""
-You are helping someone improve a scientific figure for a presentation.
-
-Figure meaning (for context):
-\"\"\"{step1_text}\"\"\"
-
-Task:
-Give 5–7 practical, specific suggestions for how to annotate this figure directly on the image
-(e.g., labels, arrows, callouts, highlights), so the key message is immediately obvious.
+You will propose how to annotate this figure for a presentation.
+
+Return ONLY valid JSON with this schema:
+{{
+  "annotations": [
+    {{
+      "type": "text" | "box" | "arrow",
+      "text": "string (only for type=text)",
+      "x": 0-1, "y": 0-1 (only for type=text),
+      "size": integer (optional, only for type=text),
+      "color": "red/blue/green/black/orange" (optional),
+
+      "xy": [x1,y1,x2,y2] (only for type=box),
+      "outline": "color" (optional for box),
+      "width": integer (optional for box/arrow),
+
+      "from": [x1,y1], "to": [x2,y2] (only for type=arrow)
+    }}
+  ]
+}}
 
 Rules:
-- Write like a human TA giving advice.
-- Avoid fancy formatting, no markdown, no bold, no bullet symbols like "-" or "•".
-- Each suggestion should be one short sentence.
-- Do not invent values; only suggest how to visually emphasize real trends.
+- Coordinates are normalized (0-1) relative to image width/height.
+- Do not invent variables not implied by the figure/paper.
+- Prefer 4–8 annotations max.
+
+Figure keywords:
+\"\"\"{kw_text}\"\"\"
 
-Return plain text, one suggestion per line.
+Paper context:
+\"\"\"{pdf_context}\"\"\"
 """
-    step2_text = call_llm(step2_prompt, model=MODEL_STEP2, temperature=0.3, max_tokens=260)
-    step2_text = clean_text(step2_text)
+    raw_step2 = ""
+    try:
+        raw_step2 = call_vision_llm(step2_prompt, image, max_tokens=380, temperature=0.3, force_json=True)
+    except Exception:
+        raw_step2 = call_text_llm(step2_prompt, max_tokens=380, temperature=0.3)
 
-    # -------- Step 3B: Plain-language explanation (replaces old paraphrase) --------
-    step3b_prompt = f"""
-Explain this same figure to a smart high-school or first-year university student.
+    ann_json = safe_json_parse(raw_step2) or {"annotations": []}
 
-Figure meaning:
-\"\"\"{step1_text}\"\"\"
+    # 渲染批注图
+    annotated_img = render_annotations(image, ann_json)
+
+    # -------- Step3: 通俗解释（按 style） ----------
+    style_map = {
+        "formal": "formal but still clear, like a polished class presentation",
+        "fluency": "smooth, natural spoken English for a talk",
+        "simple": "very simple, beginner-friendly English"
+    }
+    style_inst = style_map.get(style, style_map["fluency"])
+
+    step3_prompt = f"""
+Rewrite the explanation below into {style_inst}.
 
 Rules:
-- 4–6 sentences.
-- Use simple, conversational English.
-- Keep the science correct but avoid jargon unless necessary.
-- No numbers or stats unless they were explicitly in the meaning above.
+- Keep the scientific meaning the same.
+- Avoid sounding like an AI template.
+- 3–5 sentences.
+- No exact numeric values or p-values.
 
-Return plain text only.
+Original explanation:
+\"\"\"{step1_text}\"\"\"
 """
-    step3b_text = call_llm(step3b_prompt, model=MODEL_STEP3B, temperature=0.4, max_tokens=240)
-    step3b_text = clean_text(step3b_text)
-
-    return step1_text, step2_text, step3b_text
+    step3_text = call_text_llm(step3_prompt, max_tokens=300, temperature=0.4)
 
+    return step1_text, ann_json, annotated_img, step3_text
 
 # =========================
-# 6. Gradio UI
+# 5. Gradio UI
 # =========================
 
 with gr.Blocks() as demo:
-    gr.Markdown("## ChartSmith – AI Figure Explainer (multi-model workflow)")
+    gr.Markdown("## ChartSmith – AI Figure Explainer (JSON annotation + auto-render)")
 
     with gr.Row():
         with gr.Column():
             img_in = gr.Image(type="pil", label="Upload your scientific figure (screenshot is fine)")
-
             keywords = gr.Textbox(
                 label="Figure keywords / variables (English, comma-separated)",
                 placeholder="bleaching, coral, temperature, CO2"
             )
-
             style = gr.Radio(
                 ["formal", "fluency", "simple"],
-                value="formal",
-                label="Explanation style (for Step 1)"
+                value="fluency",
+                label="Explanation style for Step 3"
             )
-
             pdf_in = gr.File(
-                label="Upload the paper PDF (for context, optional)",
+                label="Upload the paper PDF (optional)",
                 type="filepath"
             )
-
             run_btn = gr.Button("Run workflow", variant="primary")
 
         with gr.Column():
             step1_box = gr.Textbox(
-                label="Step 1: Explanation of what the figure shows (paper-aware)",
-                lines=10
-            )
-
-            step2_box = gr.Textbox(
-                label="Step 2: Suggestions for annotating the figure",
+                label="Step 1: Explanation of what the figure shows (paper-style)",
                 lines=8
             )
-
-            step3b_box = gr.Textbox(
-                label="Step 3: Plain-language explanation (for class/presentation)",
+            step2_json = gr.JSON(
+                label="Step 2: Annotation plan (JSON)"
+            )
+            annotated_preview = gr.Image(
+                label="Step 2 Rendered: Annotated figure preview"
+            )
+            step3_box = gr.Textbox(
+                label="Step 3: Plain / presentation-friendly explanation",
                 lines=6
             )
 
     run_btn.click(
         run_workflow,
         inputs=[img_in, keywords, style, pdf_in],
-        outputs=[step1_box, step2_box, step3b_box],
+        outputs=[step1_box, step2_json, annotated_preview, step3_box]
     )
 
-
 if __name__ == "__main__":
     demo.launch()