Update app.py

app.py

@@ -1,85 +1,191 @@
 import gradio as gr
 from transformers import pipeline
 from PIL import Image, ImageEnhance
+from pypdf import PdfReader
 
+# ===================== 1. 加载模型（CPU + flan-t5-base） =====================
 
+# 图像 -> 文本说明：图像描述模型
 vision_pipe = pipeline(
     "image-to-text",
-    model="Salesforce/blip-image-captioning-base",
+    model="microsoft/git-base",
     device=-1   # CPU
 )
 
+# 文本改写 / 生成：flan-t5-base，带防重复参数
 text_pipe = pipeline(
     "text2text-generation",
     model="google/flan-t5-base",
     max_length=256,
-    device=-1   # CPU
+    num_beams=4,
+    no_repeat_ngram_size=3,   # 避免重复句子
+    device=-1                 # CPU
 )
 
 
-def analyze_and_enhance(image, style, language):
+# ===================== 2. 从论文 PDF 提取上下文 =====================
+
+def extract_paper_context(paper_file):
+    """
+    paper_file: gr.File 传进来的对象，一般是临时文件路径（str）
+    返回：截断后的论文文本字符串
+    """
+    if paper_file is None:
+        return ""
+
+    try:
+        # gradio 通常会传一个路径字符串
+        path = paper_file if isinstance(paper_file, str) else paper_file.name
+        reader = PdfReader(path)
+        text = ""
+        # 只取前几页，避免太长
+        for i, page in enumerate(reader.pages[:3]):
+            page_text = page.extract_text() or ""
+            text += page_text + "\n"
+        # 截断，防止太长
+        text = text.strip()[:4000]
+        return text
+    except Exception as e:
+        print("Error reading PDF:", e)
+        return ""
+
+
+def summarize_paper_context(raw_text):
+    """
+    用 flan-t5-base 把论文原文压缩成跟图相关的上下文总结
+    """
+    if not raw_text:
+        return ""
+
+    prompt = f"""
+You are given part of a scientific paper. 
+Summarize in 4-6 sentences the context that is most relevant for understanding one specific result figure.
+Focus on:
+- the research question or hypothesis,
+- the main variables or conditions,
+- what is being compared or measured.
+
+Paper text:
+{raw_text}
+
+Relevant context summary:
+"""
+    summary = text_pipe(prompt)[0]["generated_text"]
+    return summary
+
+
+# ===================== 3. 核心工作流函数 =====================
+
+def analyze_and_enhance(image, style, language, paper_file):
     if image is None:
         return None, "请先上传图表。", "", None, "", ""
 
+    # -------- (可选) 从论文中抽取上下文 --------
+    paper_raw = extract_paper_context(paper_file)
+    paper_context = summarize_paper_context(paper_raw) if paper_raw else ""
+
+    # -------- Step 2: Vision-LLM → 初步自动描述 --------
     raw_caption = vision_pipe(image)[0]["generated_text"]
 
-    prompt_academic = (
-        "You are helping to write the Results section of a scientific paper. "
-        "Based on the brief description below, write a more detailed academic "
-        "description of the figure in 3-4 sentences. "
-        "Explain, in generic terms, what is shown on the horizontal axis and "
-        "vertical axis (for example, 'on the horizontal axis' / 'on the vertical axis'), "
-        "describe the overall trend (increase, decrease, or differences between bars or lines), "
-        "and summarize the key comparison between groups.\n\n"
-        f"Brief description: {raw_caption}"
-    )
+    # 如果 vision 没认出是图表，就用一个安全的默认描述，避免 T5 乱写
+    rc = raw_caption.lower()
+    if not any(k in rc for k in ["bar", "chart", "graph", "plot", "curve", "line"]):
+        raw_caption = "a bar chart comparing multiple experimental groups under different conditions"
+
+    # -------- Step 3: LLM → 学术化扩写（结合论文上下文） --------
+    if paper_context:
+        context_part = f"""
+Here is the relevant context from the paper that may help you interpret the figure:
+
+{paper_context}
+"""
+    else:
+        context_part = "\n(No additional paper context is available; rely only on the figure description.)\n"
+
+    prompt_academic = f"""
+You are writing the Results section of a scientific paper.
+
+{context_part}
+
+Based on the brief description of a figure below, write a structured academic
+description of the figure in 3–4 sentences. Follow this structure:
+
+1. Identify the general type of figure (e.g., bar chart, line graph).
+2. Describe in generic terms what is shown on the horizontal axis
+   (e.g., experimental conditions, groups, time points).
+3. Describe in generic terms what is shown on the vertical axis
+   (e.g., response level, measurement value, percentage).
+4. Describe the overall trend (increase, decrease, or differences between groups).
+5. Summarize the main comparison or key pattern, making sure it is consistent
+   with the paper context above.
+
+Do NOT repeat any sentence. Do NOT invent specific numerical values or disease names.
+Keep it factual, generic, and neutral.
+
+Figure description: {raw_caption}
+"""
     academic_desc = text_pipe(prompt_academic)[0]["generated_text"]
 
-    caption_prompt = (
-        "Write a 1-2 sentence figure caption for an academic paper, "
-        "based on this description. The caption should be concise but "
-        "highlight the main comparison or trend.\n\n"
-        f"{academic_desc}\n\nCaption:"
-    )
+    # -------- Step 5: 生成图注 Caption --------
+    caption_prompt = f"""
+Write a concise 1–2 sentence figure caption for an academic paper
+based on the following description of a figure. 
+The caption should mention the type of figure and highlight the main comparison
+or trend, and it should be consistent with the paper context above.
+Do not include fabricated numbers.
+
+Description: {academic_desc}
+
+Caption:
+"""
     caption_text = text_pipe(caption_prompt)[0]["generated_text"]
 
-    summary_prompt = (
-        "Write a short 3-4 sentence paragraph that explains the key trend "
-        "and message of this figure for the Results section of a paper. "
-        "Assume the reader has not seen the figure and rely only on the "
-        "description.\n\n"
-        f"{academic_desc}\n\nParagraph:"
-    )
+    # -------- Step 5: 生成简短摘要 / 讨论 Paragraph --------
+    summary_prompt = f"""
+Write a short 3–4 sentence paragraph for the Results section of a scientific paper,
+explaining the key trend and message of this figure. Assume the reader has NOT seen
+the figure and rely only on the description below and the paper context. 
+Do not repeat sentences or invent exact numerical values.
+
+Description: {academic_desc}
+
+Paragraph:
+"""
     summary_text = text_pipe(summary_prompt)[0]["generated_text"]
 
+    # -------- 语言切换（英文 / 中文） --------
     if language == "中文":
         academic_desc_out = text_pipe(
-            f"Translate this academic description into Chinese:\n{academic_desc}"
+            f"Translate the following academic description into Chinese:\n{academic_desc}"
         )[0]["generated_text"]
         caption_text = text_pipe(
-            f"Translate this figure caption into Chinese:\n{caption_text}"
+            f"Translate the following figure caption into Chinese:\n{caption_text}"
         )[0]["generated_text"]
         summary_text = text_pipe(
-            f"Translate this paragraph into Chinese:\n{summary_text}"
+            f"Translate the following paragraph into Chinese:\n{summary_text}"
         )[0]["generated_text"]
     else:
         academic_desc_out = academic_desc
 
+    # -------- Step 4: 图像简单增强（对比度 + 亮度），让图表更清晰 --------
     img_rgb = image.convert("RGB")
     enhancer_c = ImageEnhance.Contrast(img_rgb)
-    img_c = enhancer_c.enhance(1.4)        # 对比度提高 40%
+    img_c = enhancer_c.enhance(1.4)          # 对比度 +40%
     enhancer_b = ImageEnhance.Brightness(img_c)
-    enhanced_image = enhancer_b.enhance(1.1)  # 略微提亮
+    enhanced_image = enhancer_b.enhance(1.1) # 略微提亮
 
     return image, raw_caption, academic_desc_out, enhanced_image, caption_text, summary_text
-    
+
+
+# ===================== 4. Gradio 界面 =====================
 
 with gr.Blocks() as demo:
-    gr.Markdown("# ChartSmith – AI 论文图表生成助手（CPU 版）")
+    gr.Markdown("# ChartSmith – AI 论文图表生成助手（支持上传论文上下文）")
 
     with gr.Row():
         with gr.Column():
             img_in = gr.Image(type="pil", label="上传你的学术图表（截图也可以）")
+            paper_file = gr.File(label="上传整篇论文 PDF（可选，用于更准确解释图）", file_types=[".pdf"])
             style = gr.Dropdown(
                 ["Formal academic", "Infographic", "Magazine-style"],
                 value="Formal academic",
@@ -90,14 +196,14 @@ with gr.Blocks() as demo:
         with gr.Column():
             orig_img = gr.Image(label="原始图表")
             raw_caption_box = gr.Textbox(label="Step 2: 初步自动描述（Vision-LLM）")
-            academic_box = gr.Textbox(label="Step 3: 学术化解释（Academic Explanation）")
+            academic_box = gr.Textbox(label="Step 3: 学术化解释（结合论文上下文）")
             enhanced_img = gr.Image(label="Step 4: 简单增强后的图示（对比度+亮度提升）")
             caption_box = gr.Textbox(label="Step 5: 自动生成图注（Caption）")
             summary_box = gr.Textbox(label="Step 5: 图表相关简短摘要 / 讨论")
 
     btn.click(
         analyze_and_enhance,
-        inputs=[img_in, style, language],
+        inputs=[img_in, style, language, paper_file],
         outputs=[orig_img, raw_caption_box, academic_box, enhanced_img, caption_box, summary_box]
     )