Update app.py

app.py

@@ -5,14 +5,13 @@ import torch.nn.functional as F
 import gradio as gr
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-import matplotlib.path as mpath
 import numpy as np
 from io import BytesIO
 from PIL import Image
 from transformers import AutoTokenizer, AutoModel
 
 # ==========================
-# 0. 环境初始化 (保持不变)
+# 0. 环境初始化
 # ==========================
 plt.switch_backend('Agg')
 
@@ -64,13 +63,14 @@ class ProtDualBranchEnhancedClassifier(nn.Module):
         return self.classifier_head(z_fused_gated)
 
 # ==========================
-# 2. 加载模型 (保持不变)
+# 2. 加载模型
 # ==========================
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 PLM_MODEL_NAME = "facebook/esm2_t30_150M_UR50D"
 CLASSIFIER_PATH = "best_model_esm2_t30_150M_UR50D.pth"
 LABEL_MAP_PATH = "label_map.json"
 
+# 简单检查
 if not os.path.exists(LABEL_MAP_PATH): raise FileNotFoundError(f"Missing {LABEL_MAP_PATH}")
 if not os.path.exists(CLASSIFIER_PATH): raise FileNotFoundError(f"Missing {CLASSIFIER_PATH}")
 
@@ -90,126 +90,141 @@ classifier.eval()
 print("✅ Ready.")
 
 # ==========================
-# 3. UniProt 风格绘图引擎 (修复高亮逻辑)
+# 3. 彻底修复的绘图引擎 (Centric Coordinates)
 # ==========================
 def draw_uniprot_style_cell(target_class):
     target = target_class.lower() if target_class else ""
     
-    # 配色板
+    # === 配色定义 ===
     c = {
-        'stroke': '#263238', 'cyto_bg': '#FFF9C4', 'peri_bg': '#E1F5FE',
-        'mem_line': '#546E7A', 'cw_pattern': '#90A4AE', 'dna': '#B0BEC5',
-        'ribo': '#CFD8DC', 'active': '#D32F2F', 'active_fill': '#FFCDD2'
+        'stroke': '#37474F',       # 默认深灰轮廓
+        'bg_peri': '#E1F5FE',      # 默认周质背景 (浅蓝)
+        'bg_cyto': '#FFF9C4',      # 默认胞质背景 (浅黄)
+        'highlight_stroke': '#D50000', # 高亮轮廓 (深红)
+        'highlight_fill': '#FFCDD2',   # 高亮填充 (淡红)
+        'dna': '#B0BEC5',
+        'ribo': '#90A4AE'
     }
 
-    # 状态判断
-    is_secreted = "extracellular" in target or "secreted" in target
+    # === 状态判断 ===
     is_om = "outer membrane" in target
     is_peri = "periplasm" in target
     is_cw = "cell wall" in target
     is_im = "plasma membrane" in target or "inner membrane" in target
     is_cyto = "cytoplasm" in target or "cytosol" in target
+    is_secreted = "extracellular" in target or "secreted" in target
 
+    # === 画布初始化 ===
+    # 中心点 (Cx, Cy) = (5, 3)
     fig, ax = plt.subplots(figsize=(8, 4.5), dpi=150)
     ax.set_xlim(0, 10)
     ax.set_ylim(0, 6)
     ax.axis('off')
 
-    def draw_membrane_layer(x, y, w, h, z, is_active):
-        # 修复：正确使用 is_active 判断颜色
-        color = c['active'] if is_active else c['stroke']
-        lw = 2.5 if is_active else 1.2
-        p1 = patches.FancyBboxPatch((x, y), w-h, h, boxstyle=f"round,pad=0,rounding_size={h/2}", 
-                                    fc="none", ec=color, lw=lw, zorder=z)
-        ax.add_patch(p1)
-        gap = 0.15
-        p2 = patches.FancyBboxPatch((x+gap, y+gap), w-h-2*gap, h-2*gap, boxstyle=f"round,pad=0,rounding_size={(h-2*gap)/2}", 
-                                    fc="none", ec=color, lw=lw, zorder=z)
-        ax.add_patch(p2)
-
-    # === 1. 背景层 ===
-    # Periplasm Background
-    base_capsule = patches.FancyBboxPatch((1, 1), 8, 4, boxstyle="round,pad=0,rounding_size=2", 
-                                          fc=c['peri_bg'], ec="none", zorder=1)
-    ax.add_patch(base_capsule)
+    # === 核心辅助函数：绘制绝对居中的胶囊 ===
+    def draw_centered_capsule(width, height, fill_color, edge_color, lw, z, linestyle='-'):
+        # FancyBboxPatch 的 xy 是左下角坐标。
+        # 要居中，左下角 x = CenterX - Width/2
+        x = 5.0 - width / 2
+        y = 3.0 - height / 2
+        # rounding_size 设为高度的一半，这就变成了标准的胶囊/药丸形状
+        r = height / 2 
+        
+        patch = patches.FancyBboxPatch(
+            (x, y), width, height,
+            boxstyle=f"round,pad=0,rounding_size={r}",
+            fc=fill_color, ec=edge_color, lw=lw, linestyle=linestyle, zorder=z
+        )
+        ax.add_patch(patch)
+        return x, y, width, height # 返回坐标供后续标注使用
+
+    # === 1. 绘制 Layer 1: 外膜 (Outer Membrane) ===
+    # 如果是 Periplasm 高亮，那么底色变红；否则是默认浅蓝
+    # 如果是 OuterMembrane 高亮，那么边框变红变粗
+    peri_fill = c['highlight_fill'] if is_peri else c['bg_peri']
+    om_edge = c['highlight_stroke'] if is_om else c['stroke']
+    om_lw = 3.5 if is_om else 1.5
     
-    # Cytoplasm Background
-    cyto_capsule = patches.FancyBboxPatch((1.6, 1.6), 6.8, 2.8, boxstyle="round,pad=0,rounding_size=1.4", 
-                                          fc=c['cyto_bg'], ec="none", zorder=2)
-    ax.add_patch(cyto_capsule)
-
-    # === 2. 结构层 (修复高亮) ===
+    # 绘制最大的���囊 (代表外膜轮廓 + 周质背景)
+    # 尺寸: 8.5 x 4.2
+    draw_centered_capsule(8.5, 4.2, peri_fill, om_edge, om_lw, z=1)
     
-    # A. Outer Membrane (外膜) - 最外层
-    draw_membrane_layer(1, 1, 8+2, 4, z=5, is_active=is_om)
-
-    # B. Cell Wall (细胞壁)
-    cw_color = c['active'] if is_cw else c['cw_pattern']
-    cw_lw = 2 if is_cw else 0.5
-    cw_patch = patches.FancyBboxPatch((1.3, 1.3), 7.4, 3.4, boxstyle="round,pad=0,rounding_size=1.7",
-                                      fc="none", ec=cw_color, lw=cw_lw, linestyle=(0, (5, 5)), zorder=4)
-    ax.add_patch(cw_patch)
+    # === 2. 绘制 Layer 2: 细胞壁 (Cell Wall) ===
+    # 位于中间层
+    cw_edge = c['highlight_stroke'] if is_cw else '#78909C'
+    cw_lw = 2.5 if is_cw else 1.0
+    cw_ls = '-' if is_cw else '--' # 平时虚线，高亮实线
     
-    # Periplasm Highlight Fill
-    if is_peri:
-         peri_hl = patches.FancyBboxPatch((1.1, 1.1), 7.8, 3.8, boxstyle="round,pad=0,rounding_size=1.9", 
-                                          fc=c['active_fill'], ec="none", alpha=0.5, zorder=1.5)
-         ax.add_patch(peri_hl)
-
-    # C. Inner Membrane (内膜)
-    draw_membrane_layer(1.6, 1.6, 6.8, 2.8, z=6, is_active=is_im)
-
-    # D. Cytoplasm Highlight Fill (修复：只在胞质激活时填充)
-    if is_cyto:
-        cyto_hl = patches.FancyBboxPatch((1.7, 1.7), 6.6, 2.6, boxstyle="round,pad=0,rounding_size=1.3", 
-                                         fc=c['active_fill'], ec="none", zorder=2.1)
-        ax.add_patch(cyto_hl)
-
-    # === 3. 内容细节 ===
-    # DNA
-    t = np.linspace(0, 15, 300)
-    x_dna = 5 + 2.0 * np.cos(t) * np.sin(t*0.5)
+    # 尺寸: 7.5 x 3.2
+    draw_centered_capsule(7.5, 3.2, "none", cw_edge, cw_lw, z=2, linestyle=cw_ls)
+
+    # === 3. 绘制 Layer 3: 内膜 (Inner Membrane) + 胞质 (Cytoplasm) ===
+    # 如果是 Cytoplasm 高亮，填充变红；否则默认浅黄
+    # 如果是 InnerMembrane 高亮，边框变红变粗
+    cyto_fill = c['highlight_fill'] if is_cyto else c['bg_cyto']
+    im_edge = c['highlight_stroke'] if is_im else c['stroke']
+    im_lw = 3.5 if is_im else 1.5
+
+    # 尺寸: 6.5 x 2.2
+    draw_centered_capsule(6.5, 2.2, cyto_fill, im_edge, im_lw, z=3)
+
+    # === 4. 内部细节 (DNA & Ribosomes) ===
+    # 仅装饰，画在最中心
+    # DNA 线条
+    t = np.linspace(0, 12, 200)
+    x_dna = 5 + 2.2 * np.cos(t) * np.sin(t*0.5)
     y_dna = 3 + 0.6 * np.sin(t)
-    ax.plot(x_dna, y_dna, color=c['dna'], lw=2, zorder=3, solid_capstyle='round')
+    ax.plot(x_dna, y_dna, color=c['dna'], lw=1.5, zorder=4, alpha=0.6)
 
-    # Ribosomes
+    # 核糖体 (点)
     rng = np.random.default_rng(42)
-    for _ in range(30):
-        rx, ry = rng.uniform(2.5, 7.5), rng.uniform(2.2, 3.8)
-        if not (4 < rx < 6 and 2.5 < ry < 3.5): 
-            circle = patches.Circle((rx, ry), radius=0.06, fc=c['ribo'], ec="none", zorder=3)
-            ax.add_patch(circle)
-
-    # Secreted Proteins
+    for _ in range(25):
+        # 在中心区域随机撒点
+        rx = rng.uniform(3.0, 7.0)
+        ry = rng.uniform(2.3, 3.7)
+        circle = patches.Circle((rx, ry), radius=0.05, fc=c['ribo'], zorder=4)
+        ax.add_patch(circle)
+
+    # === 5. 分泌蛋白 (Secreted) ===
     if is_secreted:
-        ax.text(5, 5.5, "Secreted", ha='center', va='bottom', color=c['active'], fontweight='bold', fontsize=11)
-        for i in range(3):
-            sx = 4 + i
-            ax.arrow(sx, 4.8, 0, 0.4, head_width=0.15, head_length=0.15, fc=c['active'], ec=c['active'], width=0.04, zorder=10)
-            circle = patches.Circle((sx, 4.7), radius=0.15, fc=c['active'], ec=c['stroke'], zorder=10)
-            ax.add_patch(circle)
-
-    # === 4. 标注 (修复指示线) ===
+        ax.text(5, 5.5, "SECRETED / EXTRACELLULAR", ha='center', va='center', 
+                color=c['highlight_stroke'], fontweight='bold')
+        # 画几个向上的箭头
+        ax.arrow(5, 5.2, 0, 0.4, head_width=0.2, fc=c['highlight_stroke'], ec=c['highlight_stroke'], width=0.05)
+
+    # === 6. 标注系统 (Labeling) ===
+    # 使用 annotate 自动画箭头指引
+    
+    # 定义各层的指引坐标 (全部取右侧中点)
+    # CenterY = 3. 
+    # OuterMembrane Edge X ≈ 5 + 8.5/2 = 9.25
+    # Periplasm X ≈ 5 + 8.0/2 = 9.0 (Inside OM)
+    # InnerMembrane Edge X ≈ 5 + 6.5/2 = 8.25
+    # Cytoplasm X ≈ 5
+    
     labels = [
-        # (Text, y_pos, x_connect, is_active)
-        ("Outer Membrane", 5, 1.0, is_om),
-        ("Periplasm", 4.2, 1.2, is_peri),
-        ("Cell Wall", 3.6, 1.3, is_cw),
-        ("Inner Membrane", 3.0, 1.6, is_im),
-        ("Cytoplasm", 2.0, 2.0, is_cyto)
+        ("Outer Membrane", (9.25, 3.0), (10, 4.5), is_om),
+        ("Periplasm",      (8.0, 3.8),  (9.5, 5.2), is_peri), # 指向胶囊上方空隙
+        ("Cell Wall",      (8.75, 3.0), (10, 3.5), is_cw),    # 指向中间虚线
+        ("Inner Membrane", (8.25, 3.0), (10, 2.5), is_im),
+        ("Cytoplasm",      (5.0, 3.0),  (5.0, 1.0), is_cyto)  # 指向中心，文字在下方
     ]
-    
-    for txt, y_anchor, x_connect, active in labels:
-        if active:
-            # 激活状态：红色加粗 + 指示线连到正确位置
-            ax.text(0.2, y_anchor, txt, ha='left', va='center', color=c['active'], fontweight='bold', fontsize=10)
-            # 修复：指示线连接到对应的 x_connect 位置
-            ax.plot([0.8, x_connect], [y_anchor, y_anchor], color=c['active'], lw=1.5, ls='-', alpha=0.8) 
-        else:
-            # 非激活状态：灰色小字
-            ax.text(0.2, y_anchor, txt, ha='left', va='center', color='#90A4AE', fontsize=8)
-
-    ax.text(5, 0.2, f"Localization: {target_class}", ha='center', va='bottom', fontsize=11, fontweight='bold', color=c['stroke'])
+
+    for txt, xy_target, xy_text, active in labels:
+        color = c['highlight_stroke'] if active else '#546E7A'
+        weight = 'bold' if active else 'normal'
+        
+        # 如果激活，画红色实线箭头；否则画灰色细箭头
+        arrow_props = dict(arrowstyle="->", color=color, lw=1.5 if active else 0.8)
+        
+        ax.annotate(txt, xy=xy_target, xytext=xy_text,
+                    arrowprops=arrow_props,
+                    fontsize=10, fontweight=weight, color=color,
+                    ha='center', va='center')
+
+    # 底部标题
+    ax.text(5, 0.2, f"Prediction: {target_class}", ha='center', va='bottom', 
+            fontsize=12, fontweight='bold', color='#263238')
 
     buf = BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight', transparent=True, dpi=150)
@@ -219,14 +234,14 @@ def draw_uniprot_style_cell(target_class):
     return img
 
 # ==========================
-# 4. 预测逻辑 (保持不变)
+# 4. 预测逻辑
 # ==========================
 def predict(sequence_input):
     if not sequence_input or sequence_input.isspace():
         raise gr.Error("Please input a protein sequence.")
     seq = "".join(sequence_input.split('\n')[1:]) if sequence_input.startswith('>') else sequence_input
     seq = re.sub(r'[^A-Z]', '', seq.upper())[:1024]
-    if not seq: raise gr.Error("Invalid Amino Acid Sequence.")
+    if not seq: raise gr.Error("Invalid Sequence.")
     
     with torch.no_grad():
         inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=1024).to(DEVICE)
@@ -238,64 +253,54 @@ def predict(sequence_input):
     top_label = idx_to_label[top_idx.item()]
     confidences = {idx_to_label[i]: float(p) for i, p in enumerate(probs)}
     
+    # 生成修正后的图
     cell_diagram = draw_uniprot_style_cell(top_label)
     return confidences, cell_diagram
 
 # ==========================
-# 5. UI 界面 (Paper-White Style, 保持不变)
+# 5. UI 界面
 # ==========================
 paper_css = """
 @import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;700&display=swap');
 body { font-family: 'Roboto', sans-serif !important; background-color: #ffffff; color: #1a1a1a; }
-.header-box { background: #ffffff; padding: 2rem 0; border-bottom: 1px solid #e5e7eb; margin-bottom: 2rem; text-align: left; }
-.header-title { font-size: 2.5rem; font-weight: 700; color: #000000; letter-spacing: -1px; line-height: 1.1; }
-.header-subtitle { font-size: 1.2rem; color: #52525b; font-weight: 300; margin-top: 10px; }
-.badge { display: inline-block; padding: 4px 10px; font-size: 0.8rem; font-weight: 500; color: #0f172a; background: #f1f5f9; border-radius: 4px; border: 1px solid #e2e8f0; margin-right: 8px; }
+.header-box { background: #ffffff; padding: 2rem 0; border-bottom: 1px solid #e5e7eb; margin-bottom: 2rem; }
+.header-title { font-size: 2.5rem; font-weight: 700; color: #000000; letter-spacing: -1px; }
+.badge { display: inline-block; padding: 4px 10px; font-size: 0.8rem; background: #f1f5f9; border: 1px solid #e2e8f0; border-radius: 4px; margin-right: 8px; }
 .content-box { background: #ffffff; border: 1px solid #e5e7eb; border-radius: 8px; padding: 1.5rem; }
-button.primary { background: #2563eb !important; color: white !important; font-weight: 500; }
 """
 
-theme = gr.themes.Base(
-    primary_hue="blue",
-    font=[gr.themes.GoogleFont("Roboto"), "ui-sans-serif", "system-ui"]
-).set(
-    body_background_fill="#ffffff",
-    block_background_fill="#ffffff",
-    block_border_width="1px",
-    block_label_background_fill="#ffffff"
+theme = gr.themes.Base(primary_hue="blue", font=[gr.themes.GoogleFont("Roboto"), "ui-sans-serif", "system-ui"]).set(
+    body_background_fill="#ffffff", block_background_fill="#ffffff", block_border_width="1px"
 )
 
 with gr.Blocks(theme=theme, css=paper_css, title="LocPred-Prok") as app:
     with gr.Column(elem_classes="header-box"):
         gr.HTML("""
             <div class="header-title">LocPred-Prok</div>
-            <div class="header-subtitle">Accurate prokaryotic subcellular localization using dual-branch protein language models</div>
-            <div style="margin-top: 15px;"><span class="badge">Article</span><span class="badge">ESM-2 Enhanced</span><span class="badge">Gram-negative</span></div>
+            <div style="font-size: 1.2rem; color: #52525b; margin: 10px 0;">Accurate prokaryotic subcellular localization using dual-branch protein language models</div>
+            <div><span class="badge">Article</span><span class="badge">ESM-2 Enhanced</span><span class="badge">Gram-negative</span></div>
         """)
 
     with gr.Tabs():
-        with gr.TabItem("Prediction Interface"):
+        with gr.TabItem("Prediction"):
             with gr.Row():
                 with gr.Column(scale=4, elem_classes="content-box"):
-                    gr.Markdown("### 1. Sequence Input")
-                    gr.Markdown("<small style='color:gray'>Enter protein sequence in FASTA format.</small>")
+                    gr.Markdown("### Sequence Input")
                     sequence_input = gr.Textbox(lines=10, show_label=False, placeholder=">Sequence...")
                     with gr.Row():
-                        clear_btn = gr.ClearButton(sequence_input, value="Clear")
+                        gr.ClearButton(sequence_input, value="Clear")
                         submit_btn = gr.Button("Analyze", variant="primary")
-                    gr.Markdown("#### Example Data")
-                    gr.Examples(examples=[[">Outer Membrane Protein A (OmpA)\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAFGGYQVNPYVGFEMGYDWLGRMPYKGSVENGAYKAQGVQLTAKLGYPITDDLDIYTRLGGMVWRADTKSNVYGKNHDTGVSPVFAGGVEYAITPEIATRLEYQWTNNIGDAHTIGTRPDNGMLSLGVSYRFGQGEAAPVVAPAPAPAPEVQTKHFTLKSDVLFNFNKATLKPEGQAALDQLYSQLSNLDPKDGSVVVLGYTDRIGSDAYNQGLSERRAQSVVDYLISKGIPADKISARGMGESNPVTGNTCDNVKQRAALIDCLAPDRRVEIEVKGIKDVVTQPQA"], [">Cytoplasmic Protein (Ribo)\nARYLGPKLKLSRREGTDLFLKSGVRAIDTKCKIEQAPGQHGARKPRLSDYGVQLREKQKVRRIYGVLERQFRNYYKEAARLKGNTGENLLALLEGRLDNVVYRMGFG"]], inputs=sequence_input, label=None)
+                    gr.Examples([
+                        [">Outer Membrane Protein\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAFGGYQVNPYVGFEMGYDWLGRMPYKGSVENGAYKAQGVQLTAKLGYPITDDLDIYTRLGGMVWRADTKSNVYGKNHDTGVSPVFAGGVEYAITPEIATRLEYQWTNNIGDAHTIGTRPDNGMLSLGVSYRFGQGEAAPVVAPAPAPAPEVQTKHFTLKSDVLFNFNKATLKPEGQAALDQLYSQLSNLDPKDGSVVVLGYTDRIGSDAYNQGLSERRAQSVVDYLISKGIPADKISARGMGESNPVTGNTCDNVKQRAALIDCLAPDRRVEIEVKGIKDVVTQPQA"],
+                        [">Cytoplasmic Protein\nARYLGPKLKLSRREGTDLFLKSGVRAIDTKCKIEQAPGQHGARKPRLSDYGVQLREKQKVRRIYGVLERQFRNYYKEAARLKGNTGENLLALLEGRLDNVVYRMGFG"]
+                    ], inputs=sequence_input, label="Examples")
 
                 with gr.Column(scale=6, elem_classes="content-box"):
-                    gr.Markdown("### 2. Localization Analysis")
-                    output_image = gr.Image(label="Schematic Visualization", show_label=False, show_download_button=True, interactive=False, type="pil", height=350)
-                    gr.Markdown("#### Probability Scores")
+                    gr.Markdown("### Localization Visualization")
+                    output_image = gr.Image(label="Visualization", show_label=False, show_download_button=True, interactive=False, type="pil", height=400)
+                    gr.Markdown("#### Confidence Scores")
                     output_label = gr.Label(num_top_classes=NUM_CLASSES, show_label=False)
 
-        with gr.TabItem("About & Citation"):
-             gr.Markdown("### Method Description\nLocPred-Prok utilizes...")
-
     submit_btn.click(fn=predict, inputs=sequence_input, outputs=[output_label, output_image])
-    clear_btn.click(lambda: [None, None], outputs=[output_label, output_image])
 
 app.launch()