Update app.py

app.py

@@ -12,9 +12,9 @@ from PIL import Image
 from transformers import AutoTokenizer, AutoModel
 
 # ==========================
-# 0. 环境初始化
+# 0. 环境初始化 (保持不变)
 # ==========================
-plt.switch_backend('Agg') # 关键：防止服务器绘图报错
+plt.switch_backend('Agg')
 
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
@@ -64,7 +64,7 @@ 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"
@@ -90,28 +90,16 @@ classifier.eval()
 print("✅ Ready.")
 
 # ==========================
-# 3. UniProt 风格绘图引擎 (Schematic Vector Style)
+# 3. UniProt 风格绘图引擎 (修复高亮逻辑)
 # ==========================
 def draw_uniprot_style_cell(target_class):
-    """
-    绘制类似 UniProt/SwissBioPics 的扁平化矢量风格细胞图。
-    特点：清晰的黑色描边，柔和的填充色，双层膜结构。
-    """
     target = target_class.lower() if target_class else ""
     
-    # === UniProt/SwissBioPics 配色板 ===
-    # 柔和、低饱和度背景，高饱和度高亮
+    # 配色板
     c = {
-        'stroke': '#263238',       # 深灰/黑色描边 (清晰感来源)
-        'cyto_bg': '#FFF9C4',      # 极淡的黄色 (胞质)
-        'peri_bg': '#E1F5FE',      # 极淡的蓝色 (周质)
-        'ext_bg': '#ffffff',       # 纯白背景
-        'mem_line': '#546E7A',     # 膜的线条颜色
-        'cw_pattern': '#90A4AE',   # 细胞壁颜色
-        'dna': '#B0BEC5',          # DNA 灰色
-        'ribo': '#CFD8DC',         # 核糖体 浅灰
-        'active': '#D32F2F',       # UniProt 经典红 (高亮)
-        'active_fill': '#FFCDD2'   # 红色填充
+        'stroke': '#263238', 'cyto_bg': '#FFF9C4', 'peri_bg': '#E1F5FE',
+        'mem_line': '#546E7A', 'cw_pattern': '#90A4AE', 'dna': '#B0BEC5',
+        'ribo': '#CFD8DC', 'active': '#D32F2F', 'active_fill': '#FFCDD2'
     }
 
     # 状态判断
@@ -122,127 +110,107 @@ def draw_uniprot_style_cell(target_class):
     is_im = "plasma membrane" in target or "inner membrane" in target
     is_cyto = "cytoplasm" in target or "cytosol" in target
 
-    # 创建画布 (宽高比 16:9)
     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. 背景层 (Backgrounds) ===
-
-    # 外膜区域背景 (Periplasm starts here)
-    # 整个胶囊的底色设为周质颜色 (Light Blue)
+    # === 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)
-
-    # 内膜区域背景 (Cytoplasm starts here)
-    # 内部胶囊底色设为胞质颜色 (Light Yellow)
+    
+    # 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. 结构层 (Structures with Strokes) ===
-
-    # A. Outer Membrane (外膜) - 双层线
+    # === 2. 结构层 (修复高亮) ===
+    
+    # A. Outer Membrane (外膜) - 最外层
     draw_membrane_layer(1, 1, 8+2, 4, z=5, is_active=is_om)
 
-    # B. Cell Wall (细胞壁) - 网格纹理层
-    # 在周质层中间画一层
+    # B. Cell Wall (细胞壁)
     cw_color = c['active'] if is_cw else c['cw_pattern']
     cw_lw = 2 if is_cw else 0.5
-    cw_alpha = 1.0 if is_cw else 0.4
     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)
     
-    # 如果是周质高亮，填充浅红色
+    # 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 (内膜) - 双层线
+    # C. Inner Membrane (内膜)
     draw_membrane_layer(1.6, 1.6, 6.8, 2.8, z=6, is_active=is_im)
 
-    # D. Cytoplasm (胞质) - 如果高亮则填充红
+    # 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. 内容细节 (Details) ===
-
-    # DNA (拟核) - 清晰的线条，不再是模糊的
+    # === 3. 内容细节 ===
+    # DNA
     t = np.linspace(0, 15, 300)
     x_dna = 5 + 2.0 * 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')
 
-    # Ribosomes (核糖体) - 规则的小圆圈，而不是随机噪点
-    # 使用规律的分布让画面更整洁
+    # 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)
-        # 避开 DNA 区域 (简单判断)
         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 (胞外)
+    # Secreted Proteins
     if is_secreted:
-        # 画几个 UniProt 风格的小球+箭头
         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
-            sy = 5.2
-            # 箭头
             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. 标注 (Annotations) ===
-    # 使用简单的线条指引
+    # === 4. 标注 (修复指示线) ===
     labels = [
-        ("Outer Membrane", 1, 5, is_om),
-        ("Periplasm", 1, 4.2, is_peri),
-        ("Cell Wall", 1, 3.6, is_cw),
-        ("Inner Membrane", 1, 3.0, is_im),
-        ("Cytoplasm", 1, 2.0, is_cyto)
+        # (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)
     ]
     
-    for txt, x_anchor, y_anchor, active in labels:
+    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)
-            # 画一条线连到细胞对应位置
-            con_x = 1.0 # 胶囊左边缘附近
-            ax.plot([0.8, con_x], [y_anchor, 3], color=c['active'], lw=1, ls='dotted', alpha=0.5) 
+            # 修复：指示线连接到对应的 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'])
 
-    # 保存
     buf = BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight', transparent=True, dpi=150)
     buf.seek(0)
@@ -251,7 +219,7 @@ def draw_uniprot_style_cell(target_class):
     return img
 
 # ==========================
-# 4. 预测逻辑
+# 4. 预测逻辑 (保持不变)
 # ==========================
 def predict(sequence_input):
     if not sequence_input or sequence_input.isspace():
@@ -274,46 +242,16 @@ def predict(sequence_input):
     return confidences, cell_diagram
 
 # ==========================
-# 5. UI 界面 (Paper-White Style)
+# 5. UI 界面 (Paper-White Style, 保持不变)
 # ==========================
-# 极简学术风 CSS：白底、黑字、细线
 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-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;
-}
-
-/* 内容卡片：无阴影，仅边框，类似论文中的 Figure */
-.content-box {
-    background: #ffffff;
-    border: 1px solid #e5e7eb;
-    border-radius: 8px;
-    padding: 1.5rem;
-}
-
-/* 按钮：学术蓝 */
+.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; }
+.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; }
 """
 
@@ -328,51 +266,29 @@ theme = gr.themes.Base(
 )
 
 with gr.Blocks(theme=theme, css=paper_css, title="LocPred-Prok") as app:
-    
-    # --- Header (Paper Style) ---
     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 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>
         """)
 
     with gr.Tabs():
         with gr.TabItem("Prediction Interface"):
             with gr.Row():
-                # Input
                 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>")
                     sequence_input = gr.Textbox(lines=10, show_label=False, placeholder=">Sequence...")
-                    
                     with gr.Row():
                         clear_btn = 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
-                    )
-
-                # Output
+                    gr.Examples(examples=[[">Outer Membrane Protein A (OmpA)\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAFGGYQVNPYVGFEMGYDWLGRMPYKGSVENGAYKAQGVQLTAKLGYPITDDLDIYTRLGGMVWRADTKSNVYGKNHDTGVSPVFAGGVEYAITPEIATRLEYQWTNNIGDAHTIGTRPDNGMLSLGVSYRFGQGEAAPVVAPAPAPAPEVQTKHFTLKSDVLFNFNKATLKPEGQAALDQLYSQLSNLDPKDGSVVVLGYTDRIGSDAYNQGLSERRAQSVVDYLISKGIPADKISARGMGESNPVTGNTCDNVKQRAALIDCLAPDRRVEIEVKGIKDVVTQPQA"], [">Cytoplasmic Protein (Ribo)\nARYLGPKLKLSRREGTDLFLKSGVRAIDTKCKIEQAPGQHGARKPRLSDYGVQLREKQKVRRIYGVLERQFRNYYKEAARLKGNTGENLLALLEGRLDNVVYRMGFG"]], inputs=sequence_input, label=None)
+
                 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")
                     output_label = gr.Label(num_top_classes=NUM_CLASSES, show_label=False)