Update app.py

app.py

@@ -5,134 +5,97 @@ from pypdf import PdfReader
 
 # ===================== 1. 加载模型（CPU + flan-t5-base） =====================
 
-# 图像 -> 文本说明：图像描述模型
 vision_pipe = pipeline(
     "image-to-text",
     model="microsoft/git-base",
-    device=-1   # CPU
+    device=-1
 )
 
-# 文本改写 / 生成：flan-t5-base，带防重复参数
 text_pipe = pipeline(
     "text2text-generation",
     model="google/flan-t5-base",
     max_length=256,
     num_beams=4,
-    no_repeat_ngram_size=3,   # 避免重复句子
-    device=-1                 # CPU
+    no_repeat_ngram_size=3,
+    device=-1
 )
 
-
-# ===================== 2. 从论文 PDF 提取上下文 =====================
+# ===================== 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 text.strip()[:4000]
+    except:
         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.
+Summarize in 4-6 sentences the research background most relevant for understanding
+a result figure. Focus on variables, comparisons, and experimental conditions.
 
 Paper text:
 {raw_text}
 
-Relevant context summary:
+Summary:
 """
-    summary = text_pipe(prompt)[0]["generated_text"]
-    return summary
+    return text_pipe(prompt)[0]["generated_text"]
 
 
-# ===================== 3. 核心工作流函数 =====================
+# ===================== 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"]
-
-    # 如果 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 → 学术化扩写（结合论文上下文） --------
+    # Academic explanation
     if paper_context:
-        context_part = f"""
-Here is the relevant context from the paper that may help you interpret the figure:
-
-{paper_context}
-"""
+        context_part = f"Here is relevant context from the paper:\n{paper_context}\n"
     else:
-        context_part = "\n(No additional paper context is available; rely only on the figure description.)\n"
+        context_part = "No additional paper context available.\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:
+Write a structured academic explanation (3–4 sentences) of the figure:
 
-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.
+1. Identify the figure type.
+2. Describe generically what is on the horizontal axis.
+3. Describe generically what is on the vertical axis.
+4. Describe overall trends.
+5. Summarize the main comparison, consistent with the paper context.
 
-Do NOT repeat any sentence. Do NOT invent specific numerical values or disease names.
-Keep it factual, generic, and neutral.
+Do NOT repeat sentences or invent specific numbers.
 
 Figure description: {raw_caption}
 """
     academic_desc = text_pipe(prompt_academic)[0]["generated_text"]
 
-    # -------- Step 5: 生成图注 Caption --------
+    # 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.
+Write a concise 1–2 sentence academic figure caption.
+Keep it consistent with the paper context.
 
 Description: {academic_desc}
 
@@ -140,12 +103,10 @@ Caption:
 """
     caption_text = text_pipe(caption_prompt)[0]["generated_text"]
 
-    # -------- Step 5: 生成简短摘要 / 讨论 Paragraph --------
+    # Summary 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.
+Write a 3–4 sentence paragraph explaining the key message of this figure,
+consistent with the paper context. Do not fabricate numbers.
 
 Description: {academic_desc}
 
@@ -153,53 +114,68 @@ Paragraph:
 """
     summary_text = text_pipe(summary_prompt)[0]["generated_text"]
 
-    # -------- 语言切换（英文 / 中文） --------
+    # Language translation
     if language == "中文":
-        academic_desc_out = text_pipe(
-            f"Translate the following academic description into Chinese:\n{academic_desc}"
-        )[0]["generated_text"]
-        caption_text = text_pipe(
-            f"Translate the following figure caption into Chinese:\n{caption_text}"
-        )[0]["generated_text"]
-        summary_text = text_pipe(
-            f"Translate the following paragraph into Chinese:\n{summary_text}"
-        )[0]["generated_text"]
-    else:
-        academic_desc_out = academic_desc
+        academic_desc = text_pipe(f"Translate into Chinese:\n{academic_desc}")[0]["generated_text"]
+        caption_text = text_pipe(f"Translate into Chinese:\n{caption_text}")[0]["generated_text"]
+        summary_text = text_pipe(f"Translate into Chinese:\n{summary_text}")[0]["generated_text"]
 
-    # -------- Step 4: 图像简单增强（对比度 + 亮度），让图表更清晰 --------
+    # 图像增强
     img_rgb = image.convert("RGB")
-    enhancer_c = ImageEnhance.Contrast(img_rgb)
-    img_c = enhancer_c.enhance(1.4)          # 对比度 +40%
-    enhancer_b = ImageEnhance.Brightness(img_c)
-    enhanced_image = enhancer_b.enhance(1.1) # 略微提亮
+    img_c = ImageEnhance.Contrast(img_rgb).enhance(1.4)
+    enhanced_image = ImageEnhance.Brightness(img_c).enhance(1.1)
 
-    return image, raw_caption, academic_desc_out, enhanced_image, caption_text, summary_text
+    return image, raw_caption, academic_desc, enhanced_image, caption_text, summary_text
 
 
-# ===================== 4. Gradio 界面 =====================
+# ===================== 4. Gradio UI（宽文本框 + 全宽布局） =====================
 
-with gr.Blocks() as demo:
-    gr.Markdown("# ChartSmith – AI 论文图表生成助手（支持上传论文上下文）")
+with gr.Blocks(css="""
+.wide_textbox textarea {
+    font-size: 15px !important;
+    line-height: 1.5 !important;
+    width: 100% !important;
+    min-height: 140px !important;
+}
+""") as demo:
+
+    gr.Markdown("# **ChartSmith – AI 论文图表生成助手（支持上传论文 PDF）**")
 
     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",
-                label="重绘风格（当前版本仅用于说明，不改变生成）"
-            )
+        with gr.Column(scale=1):
+            img_in = gr.Image(type="pil", label="上传你的学术图表")
+            paper_file = gr.File(label="上传论文 PDF（可选）", file_types=[".pdf"])
             language = gr.Radio(["English", "中文"], value="English", label="输出语言")
-            btn = gr.Button("分析并美化图表")
-        with gr.Column():
+            btn = gr.Button("分析并美化图表", variant="primary")
+
+        with gr.Column(scale=1.2):
             orig_img = gr.Image(label="原始图表")
-            raw_caption_box = gr.Textbox(label="Step 2: 初步自动描述（Vision-LLM）")
-            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: 图表相关简短摘要 / 讨论")
+
+            raw_caption_box = gr.Textbox(
+                label="Step 2: 初步自动描述（Vision-LLM）",
+                lines=4,
+                elem_classes=["wide_textbox"]
+            )
+
+            academic_box = gr.Textbox(
+                label="Step 3: 学术化解释（结合论文上下文）",
+                lines=6,
+                elem_classes=["wide_textbox"]
+            )
+
+            enhanced_img = gr.Image(label="Step 4: 增强后图示（对比度+亮度）")
+
+            caption_box = gr.Textbox(
+                label="Step 5: 自动生成图注（Caption）",
+                lines=4,
+                elem_classes=["wide_textbox"]
+            )
+
+            summary_box = gr.Textbox(
+                label="Step 5: 图表摘要 / 讨论",
+                lines=6,
+                elem_classes=["wide_textbox"]
+            )
 
     btn.click(
         analyze_and_enhance,