Update app.py

app.py

@@ -9,12 +9,10 @@ import matplotlib.pyplot as plt
 import numpy as np
 from transformers import AutoTokenizer, AutoModel
 
-# ==============================================================================
-# 0. 环境与缓存设置 (Environment Setup)
-# ==============================================================================
-# 强制使用非交互式后端，防止 matplotlib 在服务器报错
+# ==========================
+# 0. 环境与缓存
+# ==========================
 plt.switch_backend('Agg') 
-
 os.environ["HF_HOME"] = "/tmp/hf_cache"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
@@ -24,76 +22,53 @@ for path in ["/tmp/hf_cache", os.path.expanduser("~/.cache/huggingface")]:
     shutil.rmtree(path, ignore_errors=True)
     os.makedirs(path, exist_ok=True)
 
-# ==============================================================================
-# 1. 模型架构定义 (Model Architecture)
-# ==============================================================================
+# ==========================
+# 1. 模型架构 (含 Attention 输出)
+# ==========================
 class AttentionPooling(nn.Module):
     def __init__(self, d_model):
         super().__init__()
         self.attention_net = nn.Linear(d_model, 1)
 
     def forward(self, x, mask):
-        # x shape: (Batch, Seq_Len, Dim)
         attn_logits = self.attention_net(x).squeeze(2) 
         attn_logits.masked_fill_(mask == 0, -float('inf'))
         attn_weights = F.softmax(attn_logits, dim=1)
-        
-        # 返回: (Pooled_Embedding, Weights)
-        # Weights 用于 Panel D 的可视化
         return torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1), attn_weights
 
 class ProtDualBranchEnhancedClassifier(nn.Module):
     def __init__(self, d_model, projection_dim, num_classes, dropout, kernel_size):
         super().__init__()
         self.cls_projector = nn.Linear(d_model, projection_dim)
-        self.token_refiner = nn.Sequential(
-            nn.Conv1d(d_model, d_model, kernel_size, padding='same'), 
-            nn.ReLU()
-        )
+        self.token_refiner = nn.Sequential(nn.Conv1d(d_model, d_model, kernel_size, padding='same'), nn.ReLU())
         self.attention_pooling = AttentionPooling(d_model)
         self.tok_projector = nn.Linear(d_model, projection_dim)
         fused_dim = projection_dim * 2
         self.gate = nn.Sequential(nn.Linear(fused_dim, fused_dim), nn.Sigmoid())
-        self.classifier_head = nn.Sequential(
-            nn.LayerNorm(fused_dim), 
-            nn.Linear(fused_dim, fused_dim * 2), 
-            nn.ReLU(), 
-            nn.Dropout(dropout), 
-            nn.Linear(fused_dim * 2, num_classes)
-        )
+        self.classifier_head = nn.Sequential(nn.LayerNorm(fused_dim), nn.Linear(fused_dim, fused_dim * 2), nn.ReLU(), nn.Dropout(dropout), nn.Linear(fused_dim * 2, num_classes))
 
     def forward(self, cls_embedding, token_embeddings, mask):
-        # Branch 1: Global Semantic
         z_cls = self.cls_projector(cls_embedding)
-        
-        # Branch 2: Local Structural
         tok_emb_permuted = token_embeddings.permute(0, 2, 1)
         refined_tok_emb = self.token_refiner(tok_emb_permuted).permute(0, 2, 1)
-        
-        # ⚠️ 获取 Pooling 权重用于可视化
         z_tok_pooled, pooling_weights = self.attention_pooling(refined_tok_emb, mask)
         z_tok = self.tok_projector(z_tok_pooled)
-        
-        # Fusion Gate
         z_fused_concat = torch.cat([z_cls, z_tok], dim=1)
         gate_values = self.gate(z_fused_concat)
         z_fused_gated = z_fused_concat * gate_values
-        
         return self.classifier_head(z_fused_gated), pooling_weights
 
-# ==============================================================================
-# 2. 加载模型与配置 (Load Resources)
-# ==============================================================================
+# ==========================
+# 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}")
 
-# 加载 Label Map
 with open(LABEL_MAP_PATH, 'r') as f:
     label_to_idx = json.load(f)
     idx_to_label = {v: k for k, v in label_to_idx.items()}
@@ -103,41 +78,36 @@ D_MODEL = 640
 print("🔹 Loading models...")
 tokenizer = AutoTokenizer.from_pretrained(PLM_MODEL_NAME)
 plm_model = AutoModel.from_pretrained(PLM_MODEL_NAME).to(DEVICE).eval()
-
 classifier = ProtDualBranchEnhancedClassifier(D_MODEL, 32, NUM_CLASSES, 0.3, 3).to(DEVICE)
 classifier.load_state_dict(torch.load(CLASSIFIER_PATH, map_location=DEVICE))
 classifier.eval()
 print("✅ Ready.")
 
-# ==============================================================================
-# 3. Panel B: SVG 绘图引擎 (Visualization Engine)
-# ==============================================================================
+# ==========================
+# 3. Panel B: SVG 细胞图 (6类完整显示)
+# ==========================
 def generate_bacterial_svg(target_class):
     target = target_class.lower() if target_class else ""
     
-    # 1. 状态判断
-    is_om = "outer membrane" in target
+    # 状态
+    is_sec  = "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_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
 
-    # 2. 颜色配置 (高对比度科研风)
+    # 颜色
     c = {
-        # 激活态: 鲜红
         "hl_stroke": "#D32F2F", "hl_fill": "#FFEBEE", "hl_text": "#B71C1C", "hl_dot": "#D32F2F",
-        # 未激活态: 极淡的灰白 (背景化)
         "bg_stroke": "#90A4AE", "bg_fill_om": "#F5F5F5", "bg_fill_im": "#FAFAFA",
         "bg_text": "#78909C", "bg_line": "#CFD8DC", "bg_dot": "#B0BEC5"
     }
 
-    # 3. 样式生成器 (这里修复了之前的 bug)
+    # 结构样式
     def style(active, base_fill, base_stroke, w_act="4", w_norm="2"):
-        if active: 
-            return c["hl_fill"], c["hl_stroke"], w_act
-        # ✅ 修复点：这里原来写成了 width_norm，现已修正为 w_norm
-        return base_fill, base_stroke, w_norm
+        if active: return c["hl_fill"], c["hl_stroke"], w_act
+        return base_fill, base_stroke, width_norm
 
     om_f, om_s, om_w = style(is_peri, c["bg_fill_om"], c["hl_stroke"] if is_om else c["bg_stroke"])
     cw_s = c["hl_stroke"] if is_cw else "#B0BEC5"
@@ -149,37 +119,47 @@ def generate_bacterial_svg(target_class):
         if active: return c["hl_text"], "bold", c["hl_stroke"], "2.5", c["hl_dot"], "5"
         return c["bg_text"], "normal", c["bg_line"], "1.5", c["bg_dot"], "3"
 
-    l_om, l_peri, l_cw, l_im, l_cyto = label_style(is_om), label_style(is_peri), label_style(is_cw), label_style(is_im), label_style(is_cyto)
+    l_sec  = label_style(is_sec)
+    l_om   = label_style(is_om)
+    l_peri = label_style(is_peri)
+    l_cw   = label_style(is_cw)
+    l_im   = label_style(is_im)
+    l_cyto = label_style(is_cyto)
 
-    # 4. 坐标定义
-    bx, by = 280, 210 # 细菌中心
-    tx = 600          # 标签文字起始 X 坐标
+    # 坐标系统 (中心 280, 210)
+    bx, by = 280, 210
+    tx = 600 # 标签起始X
 
+    # 锚点目标 (Target Anchor Points)
     targets = {
-        "om":   (bx + 140, by - 120), # 外膜线
-        "peri": (bx + 120, by - 90),  # 周质间隙
-        "cw":   (bx + 100, by - 70),  # 细胞壁线
-        "im":   (bx + 70,  by - 50),  # 内膜线
-        "cyto": (bx,       by)        # 胞质中心
+        "sec":  (bx,       by - 180), # 胞外 (悬浮在上方)
+        "om":   (bx + 140, by - 120), # 外膜
+        "peri": (bx + 120, by - 90),  # 周质
+        "cw":   (bx + 100, by - 70),  # 细胞壁
+        "im":   (bx + 70,  by - 50),  # 内膜
+        "cyto": (bx,       by)        # 胞质
     }
     
-    text_y = {"om": 90, "peri": 150, "cw": 210, "im": 270, "cyto": 330}
+    # 标签文字Y坐标 (均匀分布6个)
+    text_y = {
+        "sec": 50, "om": 110, "peri": 170, "cw": 230, "im": 290, "cyto": 350
+    }
 
-    # 5. 贝塞尔曲线连接器
+    # 贝塞尔曲线生成器
     def draw_connector(key, style_tuple, label_text):
         txt_col, weight, line_col, width, dot_col, r = style_tuple
         tx_pos, ty_pos = tx, text_y[key]
         ex, ey = targets[key]
         
-        # 贝塞尔控制点
-        c1x, c1y = tx_pos - 100, ty_pos 
-        c2x, c2y = ex + 50, ey
+        # S形曲线控制点
+        c1x, c1y = tx_pos - 80, ty_pos 
+        c2x, c2y = ex + 60, ey
         
         path = f"M {tx_pos - 10} {ty_pos - 5} C {c1x} {c1y}, {c2x} {c2y}, {ex} {ey}"
         
         return f"""
         <g>
-            <text x="{tx_pos}" y="{ty_pos}" fill="{txt_col}" font-weight="{weight}" font-size="15" font-family="Arial">{label_text}</text>
+            <text x="{tx_pos}" y="{ty_pos}" fill="{txt_col}" font-weight="{weight}" font-size="14" font-family="Arial">{label_text}</text>
             <path d="{path}" fill="none" stroke="{line_col}" stroke-width="{width}" />
             <circle cx="{ex}" cy="{ey}" r="{r}" fill="{dot_col}" stroke="white" stroke-width="1" />
         </g>
@@ -196,15 +176,7 @@ def generate_bacterial_svg(target_class):
             </g>
         </g>
 
-        {f'''
-        <g transform="translate(500, 40)">
-            <text x="0" y="0" text-anchor="middle" fill="{c['hl_stroke']}" font-weight="bold" font-family="Arial" font-size="14">SECRETED</text>
-            <path d="M 0 10 L 0 40" stroke="{c['hl_stroke']}" stroke-width="2" marker-end="url(#arrow_hl)" />
-        </g>
-        ''' if is_secreted else ""}
-        
-        <defs><marker id="arrow_hl" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto"><polygon points="0 0, 10 3.5, 0 7" fill="{c['hl_stroke']}" /></marker></defs>
-
+        {draw_connector("sec", l_sec, "Extracellular / Secreted")}
         {draw_connector("om", l_om, "Outer Membrane")}
         {draw_connector("peri", l_peri, "Periplasm")}
         {draw_connector("cw", l_cw, "Cell Wall")}
@@ -213,46 +185,48 @@ def generate_bacterial_svg(target_class):
     </svg>"""
     return svg
 
-# ==============================================================================
-# 4. Panel D: Attention 绘图引擎 (Interpretability)
-# ==============================================================================
-def draw_pooling_weights(weights, sequence):
+# ==========================
+# 4. Panel D: Attention Heatmap (热图版)
+# ==========================
+def draw_attention_heatmap_strip(weights, sequence):
     """
-    Visualize Attention Pooling Weights (1D Heatmap/Bar).
+    Draws a 1D Heatmap Strip for Attention Weights.
+    Standard Bioinformatics visualization style.
     """
-    # 归一化
+    # 归一化 (0-1)
     if weights.max() > 0:
         weights = (weights - weights.min()) / (weights.max() - weights.min())
+    
+    # 准备数据 (Reshape to 2D for imshow: [1, Seq_Len])
+    data = weights.reshape(1, -1)
 
-    fig, ax = plt.subplots(figsize=(6, 3), dpi=120)
-    x = np.arange(len(weights))
+    fig, ax = plt.subplots(figsize=(8, 1.5), dpi=150) # 长条形
     
-    # 绘制红色条形
-    ax.bar(x, weights, width=1.0, color='#D32F2F', alpha=0.8, label='Attention')
+    # 绘制热图 (使用 Reds 色系，颜色越深 Attention 越高)
+    im = ax.imshow(data, cmap='Reds', aspect='auto', vmin=0, vmax=1)
     
-    # 样式
-    ax.set_title("Learned Motif Importance (Attention Pooling)", fontsize=10, fontweight='bold', color='#37474F')
+    # 样式美化
+    ax.set_title("Sequence Attention Heatmap (High Color = Key Motif)", fontsize=10, fontweight='bold', color='#37474F', pad=10)
     ax.set_xlabel("Residue Position", fontsize=9)
-    ax.set_ylabel("Weight", fontsize=9)
-    ax.spines['top'].set_visible(False)
-    ax.spines['right'].set_visible(False)
-    ax.spines['left'].set_visible(False)
-    ax.set_yticks([]) 
     
-    # 标注最高峰 (Key Motif)
-    threshold = np.percentile(weights, 98)
-    if weights.max() > threshold:
-        peak_idx = np.argmax(weights)
-        ax.annotate('Key Motif', xy=(peak_idx, weights[peak_idx]), xytext=(peak_idx, weights[peak_idx]+0.2),
-                    arrowprops=dict(facecolor='#37474F', shrink=0.05, width=1, headwidth=5),
-                    ha='center', fontsize=8, color='#37474F')
+    # 隐藏 Y 轴刻度
+    ax.set_yticks([])
+    
+    # 添加 Colorbar
+    cbar = plt.colorbar(im, ax=ax, orientation='vertical', fraction=0.02, pad=0.02)
+    cbar.ax.tick_params(labelsize=8)
+    cbar.outline.set_visible(False)
+
+    # 隐藏边框
+    for spine in ax.spines.values():
+        spine.set_visible(False)
 
     plt.tight_layout()
     return fig
 
-# ==============================================================================
-# 5. 预测主逻辑 (Prediction Logic)
-# ==============================================================================
+# ==========================
+# 5. 预测主逻辑
+# ==========================
 def predict(sequence_input):
     if not sequence_input or sequence_input.isspace(): raise gr.Error("Empty Input")
     
@@ -266,35 +240,33 @@ def predict(sequence_input):
         
         hidden_states = outputs.last_hidden_state
         cls_embedding = hidden_states[:, 0, :]
-        token_embeddings = hidden_states[:, 1:-1, :] # No CLS/EOS
+        token_embeddings = hidden_states[:, 1:-1, :]
         token_mask = inputs['attention_mask'][:, 1:-1]
         
-        # ⚠️ 获取 logits 和 pooling_weights
         logits, pooling_weights = classifier(cls_embedding, token_embeddings, token_mask)
         probs = F.softmax(logits, dim=1)[0]
     
-    # 1. 结果 (Panel C)
+    # 1. 结果
     top_label = idx_to_label[torch.max(probs, dim=0)[1].item()]
     confidences = {idx_to_label[i]: float(p) for i, p in enumerate(probs)}
     
-    # 2. Panel B: SVG
+    # 2. SVG (Panel B)
     svg = generate_bacterial_svg(top_label)
     
-    # 3. Panel D: Attention Plot
-    # 取 batch 中第一个样本的 weights
+    # 3. Heatmap (Panel D)
     w_np = pooling_weights[0].cpu().numpy()
-    attn_plot = draw_pooling_weights(w_np, seq)
+    heatmap_plot = draw_attention_heatmap_strip(w_np, seq)
     
-    return confidences, svg, attn_plot
+    return confidences, svg, heatmap_plot
 
-# ==============================================================================
-# 6. UI 布局 (Four-Block Paper Style)
-# ==============================================================================
+# ==========================
+# 6. UI Layout (4-Block)
+# ==========================
 layout_css = """
 @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;800&display=swap');
 body { background-color: #ffffff; font-family: 'Inter', sans-serif; }
 
-/* Header: Sky Blue Theme */
+/* Header */
 .header-div {
     background: linear-gradient(to right, #E0F7FA, #E1F5FE);
     padding: 1.5rem;
@@ -330,7 +302,6 @@ theme = gr.themes.Soft(primary_hue="sky").set(body_background_fill="white", bloc
 
 with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok") as app:
     
-    # --- Header ---
     gr.HTML("""
         <div class="header-div">
             <div class="header-title">LocPred-Prok</div>
@@ -338,7 +309,7 @@ with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok") as app:
         </div>
     """)
 
-    # --- Row 1: Panels A & B ---
+    # Row 1
     with gr.Row():
         with gr.Column(elem_classes="panel-card"):
             gr.Markdown("<div class='panel-header'><span class='panel-label'>A</span>Sequence Input</div>")
@@ -354,14 +325,14 @@ with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok") as app:
             gr.Markdown("<div class='panel-header'><span class='panel-label'>B</span>Localization Visualization</div>")
             output_svg = gr.HTML(label="Visual", show_label=False)
 
-    # --- Row 2: Panels C & D ---
+    # Row 2
     with gr.Row():
         with gr.Column(elem_classes="panel-card"):
             gr.Markdown("<div class='panel-header'><span class='panel-label'>C</span>Prediction Confidence</div>")
             output_label = gr.Label(num_top_classes=NUM_CLASSES, show_label=False)
 
         with gr.Column(elem_classes="panel-card"):
-            gr.Markdown("<div class='panel-header'><span class='panel-label'>D</span>Learned Motif Importance (Attention)</div>")
+            gr.Markdown("<div class='panel-header'><span class='panel-label'>D</span>Attention Heatmap (Motif Discovery)</div>")
             output_plot = gr.Plot(label="Attention", show_label=False)
 
     submit_btn.click(fn=predict, inputs=sequence_input, outputs=[output_label, output_svg, output_plot])