Update app.py

app.py

@@ -85,64 +85,36 @@ classifier.eval()
 print("✅ Ready.")
 
 # ==========================
-# 3. 标签标准化映射 (核心修正)
+# 3. 标签标准化映射
 # ==========================
 def clean_label_name(raw_label):
-    """
-    将模型原始输出的各类写法（如 OuterMembrane, Cytoplasmic）
-    统一映射为您要求的 6 个标准显示名称。
-    """
     raw = raw_label.strip()
-    
-    # 映射字典：Key 是模型可能的输出，Value 是标准显示名称
     mapping = {
-        # 1. Outer membrane
         "OuterMembrane": "Outer membrane", "Outer membrane": "Outer membrane",
-        
-        # 2. Periplasm
         "Periplasmic": "Periplasm", "Periplasm": "Periplasm",
-        
-        # 3. Cell wall
         "Cellwall": "Cell wall", "Cell wall": "Cell wall",
-        
-        # 4. Cytoplasmic membrane (即 Inner Membrane)
-        "CYtoplasmicMembrane": "Cytoplasmic membrane", "Cytoplasmic membrane": "Cytoplasmic membrane",
-        "InnerMembrane": "Cytoplasmic membrane", "Inner membrane": "Cytoplasmic membrane",
-        
-        # 5. Cytoplasm
+        "CYtoplasmicMembrane": "Cytoplasmic membrane", "InnerMembrane": "Cytoplasmic membrane",
         "Cytoplasmic": "Cytoplasm", "Cytoplasm": "Cytoplasm",
-        
-        # 6. Extracellular
         "Extracellular": "Extracellular", "Secreted": "Extracellular"
     }
-    
-    # 尝试直接匹配
-    if raw in mapping:
-        return mapping[raw]
-    
-    # 尝试忽略大小写匹配
+    if raw in mapping: return mapping[raw]
     raw_lower = raw.lower()
     for k, v in mapping.items():
-        if k.lower() == raw_lower:
-            return v
-            
-    return raw # 如果没匹配上，返回原样
+        if k.lower() == raw_lower: return v
+    return raw
 
 # ==========================
-# 4. SVG 引擎 (适配标准标签)
+# 4. SVG 引擎 (纯净展示版 - 无下载按钮)
 # ==========================
 def infer_gram_type(std_label):
-    """基于标准标签推断"""
     if std_label in ["Outer membrane", "Periplasm"]: return "negative"
     if std_label == "Cell wall": return "positive"
     return "negative"
 
 def generate_scientific_svg(target_class):
-    # 先转为标准标签
     std_target = clean_label_name(target_class)
     gram_type = infer_gram_type(std_target)
     
-    # 状态判断 (使用标准名称)
     is_sec  = (std_target == "Extracellular")
     is_om   = (std_target == "Outer membrane")
     is_peri = (std_target == "Periplasm")
@@ -160,22 +132,18 @@ def generate_scientific_svg(target_class):
     cx, cy = 300, 210
     tx = 620
 
-    # 绘制主体
     shapes = ""
     if gram_type == 'negative':
-        # Outer Membrane
         col_om = c['hl_stroke'] if is_om else c['bg_stroke']
         fill_om = c['hl_fill'] if is_peri else c['bg_fill']
         w_om = "4" if is_om else "2"
         shapes += f'<rect x="{cx-200}" y="{cy-120}" width="400" height="240" rx="120" ry="120" fill="{fill_om}" stroke="{col_om}" stroke-width="{w_om}" />'
         
-        # Cell Wall
         col_cw = c['hl_stroke'] if is_cw else '#B0BEC5'
         w_cw = "3" if is_cw else "1.5"
         dash_cw = "0" if is_cw else "6,4"
         shapes += f'<rect x="{cx-170}" y="{cy-90}" width="340" height="180" rx="90" ry="90" fill="none" stroke="{col_cw}" stroke-width="{w_cw}" stroke-dasharray="{dash_cw}" />'
         
-        # Cytoplasmic Membrane (Inner)
         col_im = c['hl_stroke'] if is_im else c['bg_stroke']
         fill_im = c['hl_fill'] if is_cyto else c['bg_fill']
         w_im = "4" if is_im else "2"
@@ -186,7 +154,6 @@ def generate_scientific_svg(target_class):
             "cw": (cx+170, cy), "im": (cx+140, cy+30), "cyto": (cx, cy)
         }
     else:
-        # Gram Positive (Thick Wall)
         col_cw = c['hl_stroke'] if is_cw else c['bg_stroke']
         fill_bg = c['hl_fill'] if is_peri else c['bg_fill']
         w_cw = "6" if is_cw else "4"
@@ -215,7 +182,6 @@ def generate_scientific_svg(target_class):
             <path d="M 0 5 L 0 30" stroke="{c['hl_stroke']}" stroke-width="2" marker-end="url(#arrow_hl)" />
         </g>"""
 
-    # --- 标签列表 (使用您要求的标准名称) ---
     labels_config = [
         ("Extracellular", "sec", is_sec),
         ("Outer membrane", "om", is_om),
@@ -266,88 +232,49 @@ def generate_scientific_svg(target_class):
         <text x="400" y="400" text-anchor="middle" font-family="'Lato', sans-serif" font-size="16" fill="#546E7A" font-weight="bold">Prediction: {std_target}</text>
     </svg>"""
     
-    html = f"""<div>{final_svg}
-    <div style="display:flex; justify-content:center; gap:10px; margin-top:5px;">
-        <button onclick="downloadSVG('{svg_id}')" style="font-size:11px; padding:4px 8px; border:1px solid #ccc; border-radius:4px; cursor:pointer; font-family:'Lato', sans-serif;">Download SVG</button>
-    </div>
-    <script>
-        function downloadSVG(id) {{
-            const svg = document.getElementById(id);
-            const s = new XMLSerializer().serializeToString(svg);
-            const b = new Blob([s], {{type: "image/svg+xml;charset=utf-8"}});
-            const u = URL.createObjectURL(b);
-            const a = document.createElement("a"); a.href = u; a.download = "cell_loc.svg";
-            document.body.appendChild(a); a.click(); document.body.removeChild(a);
-        }}
-    </script></div>"""
+    # 纯净 HTML，无按钮
+    html = f"<div style='text-align:center;'>{final_svg}</div>"
     return html
 
 # ==========================
-# 4. Panel D: Wrapped Attention Heatmap (折行显示氨基酸)
+# 4. Wrapped Attention Heatmap
 # ==========================
 def draw_wrapped_attention_heatmap(weights, sequence, chars_per_line=60):
-    """
-    绘制折行热图：每行显示固定数量的氨基酸，下方显示字母。
-    """
-    # 归一化权重 (0-1)
-    if weights.max() > 0:
-        weights = (weights - weights.min()) / (weights.max() - weights.min())
-    
+    if weights.max() > 0: weights = (weights - weights.min()) / (weights.max() - weights.min())
     seq_len = len(sequence)
-    # 计算行数
     num_rows = (seq_len + chars_per_line - 1) // chars_per_line
-    
-    # 动态调整画布高度
     fig_height = max(2, num_rows * 0.8) 
     fig, axes = plt.subplots(num_rows, 1, figsize=(10, fig_height), dpi=150)
-    
-    # 如果只有一行，axes不是列表，强制转列表
-    if num_rows == 1:
-        axes = [axes]
-        
-    # 字体设置
     plt.rcParams['font.family'] = 'sans-serif'
     plt.rcParams['font.sans-serif'] = ['Lato', 'monospace']
+    if num_rows == 1: axes = [axes]
 
     for i in range(num_rows):
         ax = axes[i]
         start_idx = i * chars_per_line
         end_idx = min((i + 1) * chars_per_line, seq_len)
-        
-        # 截取片段
         sub_weights = weights[start_idx:end_idx]
         sub_seq = sequence[start_idx:end_idx]
         current_len = len(sub_seq)
         
-        # 补全最后一行以便绘图 (保持对齐)
         display_weights = np.zeros((1, chars_per_line))
         display_weights[0, :current_len] = sub_weights
         
-        # 绘制热图 (Reds)
         im = ax.imshow(display_weights, cmap='Reds', aspect='auto', vmin=0, vmax=1)
         
-        # 在每个格子上写氨基酸字母
         for j, char in enumerate(sub_seq):
-            # 文字颜色根据背景深浅调整 (这里简化为黑色，因为权重一般不会全黑)
             ax.text(j, 0, char, ha='center', va='center', fontsize=9, color='black', fontweight='bold')
             
-        # 设置 X 轴 (不显示刻度，只显示格子边界)
         ax.set_xticks(np.arange(chars_per_line) - 0.5, minor=True)
         ax.set_yticks([])
         ax.grid(which="minor", color="w", linestyle='-', linewidth=1)
         ax.tick_params(which="minor", bottom=False, left=False)
         ax.tick_params(which="major", bottom=False, left=False, labelbottom=False)
-        
-        # 隐藏边框
-        for spine in ax.spines.values():
-            spine.set_visible(False)
-            
-        # 左侧添加行号索引 (1, 61, 121...)
+        for spine in ax.spines.values(): spine.set_visible(False)
         ax.set_ylabel(f"{start_idx+1}", rotation=0, ha='right', va='center', fontsize=10, color='#546E7A')
 
     plt.tight_layout()
-    # 顶部标题
-    fig.suptitle(f"Attention Heatmap with Sequence ({seq_len} residues)", fontsize=12, fontweight='bold', color='#37474F', y=1.02)
+    fig.suptitle(f"Attention Heatmap (Sequence Length: {seq_len})", fontsize=12, fontweight='bold', color='#37474F', y=1.02)
     return fig
 
 # ==========================
@@ -364,39 +291,50 @@ def predict(sequence_input):
         logits, pooling_weights = classifier(outputs.last_hidden_state[:, 0, :], outputs.last_hidden_state[:, 1:-1, :], inputs['attention_mask'][:, 1:-1])
         probs = F.softmax(logits, dim=1)[0]
     
-    # 获取原始预测 ID
     top_id = torch.max(probs, dim=0)[1].item()
-    # 获取原始标签文本 (如 OuterMembrane)
     raw_label = idx_to_label[top_id]
-    
-    # 转换为标准显示文本 (如 Outer membrane)
     clean_top_label = clean_label_name(raw_label)
     
-    # 构建置信度字典 (全部转换为标准名称)
     confidences = {}
     for i, p in enumerate(probs):
         orig_name = idx_to_label[i]
         std_name = clean_label_name(orig_name)
         confidences[std_name] = float(p)
     
-    # 绘图 (传入标准名称)
     svg = generate_scientific_svg(clean_top_label)
-    
-    # 绘制折行热图
     heatmap = draw_wrapped_attention_heatmap(pooling_weights[0].cpu().numpy(), seq, chars_per_line=60)
     
     return confidences, svg, heatmap
 
 # ==========================
-# 6. UI Layout (UniProt Font Style)
+# 6. UI Layout (Enhanced Header)
 # ==========================
 layout_css = """
 @import url('https://fonts.googleapis.com/css2?family=Lato:wght@300;400;700&display=swap');
 body, button, input, textarea, .gradio-container { font-family: 'Lato', sans-serif !important; }
 
-.header-div { background: linear-gradient(to right, #E0F7FA, #E1F5FE); padding: 1.5rem; border-radius: 8px; margin-bottom: 20px; text-align: center; border: 1px solid #B3E5FC; }
+/* Header 样式 */
+.header-div { 
+    background: linear-gradient(to right, #E0F7FA, #E1F5FE); 
+    padding: 1.5rem; border-radius: 8px; margin-bottom: 20px; 
+    text-align: center; border: 1px solid #B3E5FC; 
+}
 .header-title { font-size: 2.2rem; font-weight: 800; color: #0288D1; margin-bottom: 5px; }
-.header-sub { font-size: 1.0rem; color: #0277BD; }
+.header-sub { font-size: 1.0rem; color: #0277BD; margin-bottom: 12px; }
+
+/* Badge 链接样式 */
+.badge-container { display: flex; justify-content: center; gap: 10px; flex-wrap: wrap; }
+.badge-link {
+    text-decoration: none; display: inline-flex; align-items: center;
+    background-color: #ffffff; color: #334155; 
+    padding: 4px 10px; border-radius: 6px; 
+    font-size: 0.85rem; font-weight: 600;
+    border: 1px solid #cbd5e1; transition: all 0.2s;
+}
+.badge-link:hover { background-color: #f1f5f9; border-color: #0288D1; color: #0288D1; }
+.badge-icon { margin-right: 5px; }
+
+/* Panel 样式 */
 .panel-card { border: 1px solid #e2e8f0; border-radius: 8px; padding: 15px; background: white; height: 100%; display: flex; flex-direction: column; }
 .panel-header { font-weight: 700; color: #475569; border-bottom: 2px solid #f1f5f9; padding-bottom: 8px; margin-bottom: 12px; font-size: 1.0rem; }
 .panel-label { display: inline-block; background: #E0F7FA; color: #0277BD; border: 1px solid #B2EBF2; padding: 2px 8px; border-radius: 4px; font-size: 0.8rem; margin-right: 8px; font-weight: 800; }
@@ -405,7 +343,28 @@ body, button, input, textarea, .gradio-container { font-family: 'Lato', sans-ser
 theme = gr.themes.Soft(primary_hue="sky").set(body_background_fill="white", block_background_fill="white", block_border_width="0px")
 
 with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok") as app:
-    gr.HTML("""<div class="header-div"><div class="header-title">LocPred-Prok</div><div class="header-sub">Deep Learning Framework for Prokaryotic Subcellular Localization</div></div>""")
+    
+    # --- Enhanced Header ---
+    gr.HTML("""
+        <div class="header-div">
+            <div class="header-title">LocPred-Prok</div>
+            <div class="header-sub">Dual-Branch Deep Learning for Prokaryotic Subcellular Localization</div>
+            <div class="badge-container">
+                <a href="https://github.com/isyslab-hust/LocPred-Prok" target="_blank" class="badge-link">
+                    GitHub
+                </a>
+                <a href="#" target="_blank" class="badge-link">
+                    Paper
+                </a>
+                <span class="badge-link" style="cursor:default">
+                    🧬 ESM-2 Enhanced
+                </span>
+                <span class="badge-link" style="cursor:default">
+                    ⚖️ MIT License
+                </span>
+            </div>
+        </div>
+    """)
 
     with gr.Row():
         with gr.Column(elem_classes="panel-card"):