Update app.py

app.py

@@ -85,41 +85,70 @@ classifier.eval()
 print("✅ Ready.")
 
 # ==========================
-# 3. 标签映射与 SVG 引擎 (UniProt Font)
+# 3. 标签标准化映射 (核心修正)
 # ==========================
-
-def map_model_label_to_svg_key(model_label):
-    l = model_label.strip()
+def clean_label_name(raw_label):
+    """
+    将模型原始输出的各类写法（如 OuterMembrane, Cytoplasmic）
+    统一映射为您要求的 6 个标准显示名称。
+    """
+    raw = raw_label.strip()
+    
+    # 映射字典：Key 是模型可能的输出，Value 是标准显示名称
     mapping = {
-        "Extracellular": "sec",
-        "OuterMembrane": "om",
-        "Periplasmic": "peri",
-        "Cellwall": "cw",
-        "CYtoplasmicMembrane": "im",
-        "Cytoplasmic": "cyto"
+        # 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
+        "Cytoplasmic": "Cytoplasm", "Cytoplasm": "Cytoplasm",
+        
+        # 6. Extracellular
+        "Extracellular": "Extracellular", "Secreted": "Extracellular"
     }
-    if l in mapping: return mapping[l]
-    l_lower = l.lower()
+    
+    # 尝试直接匹配
+    if raw in mapping:
+        return mapping[raw]
+    
+    # 尝试忽略大小写匹配
+    raw_lower = raw.lower()
     for k, v in mapping.items():
-        if k.lower() == l_lower: return v
-    return None
+        if k.lower() == raw_lower:
+            return v
+            
+    return raw # 如果没匹配上，返回原样
 
-def infer_gram_type(model_label):
-    key = map_model_label_to_svg_key(model_label)
-    if key in ["om", "peri"]: return "negative"
-    if key == "cw": return "positive"
+# ==========================
+# 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):
-    active_key = map_model_label_to_svg_key(target_class)
-    gram_type = infer_gram_type(target_class)
+    # 先转为标准标签
+    std_target = clean_label_name(target_class)
+    gram_type = infer_gram_type(std_target)
     
-    is_sec  = (active_key == "sec")
-    is_om   = (active_key == "om")
-    is_peri = (active_key == "peri")
-    is_cw   = (active_key == "cw")
-    is_im   = (active_key == "im")
-    is_cyto = (active_key == "cyto")
+    # 状态判断 (使用标准名称)
+    is_sec  = (std_target == "Extracellular")
+    is_om   = (std_target == "Outer membrane")
+    is_peri = (std_target == "Periplasm")
+    is_cw   = (std_target == "Cell wall")
+    is_im   = (std_target == "Cytoplasmic membrane")
+    is_cyto = (std_target == "Cytoplasm")
 
     c = {
         'hl_stroke': '#D32F2F', 'hl_fill': '#FFEBEE', 'hl_text': '#B71C1C', 'hl_dot': '#D32F2F',
@@ -131,19 +160,22 @@ 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"
@@ -154,6 +186,7 @@ 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"
@@ -178,16 +211,17 @@ def generate_scientific_svg(target_class):
     sec_svg = ""
     if is_sec:
         sec_svg = f"""<g transform="translate({cx}, {cy-170})">
-            <text x="0" y="0" text-anchor="middle" fill="{c['hl_text']}" font-weight="bold" font-family="'Lato', 'Helvetica Neue', Arial, sans-serif" font-size="14">SECRETED</text>
+            <text x="0" y="0" text-anchor="middle" fill="{c['hl_text']}" font-weight="bold" font-family="'Lato', sans-serif" font-size="14">EXTRACELLULAR</text>
             <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),
+        ("Outer membrane", "om", is_om),
         ("Periplasm", "peri", is_peri),
-        ("Cell Wall", "cw", is_cw),
-        ("Inner Membrane", "im", is_im),
+        ("Cell wall", "cw", is_cw),
+        ("Cytoplasmic membrane", "im", is_im),
         ("Cytoplasm", "cyto", is_cyto)
     ]
     if gram_type == 'positive':
@@ -214,7 +248,7 @@ def generate_scientific_svg(target_class):
         
         label_svg += f"""
         <g>
-            <text x="{tx}" y="{ty}" fill="{col_txt}" font-weight="{w_txt}" font-size="14" font-family="'Lato', 'Helvetica Neue', Arial, sans-serif">{text}</text>
+            <text x="{tx}" y="{ty}" fill="{col_txt}" font-weight="{w_txt}" font-size="14" font-family="'Lato', sans-serif">{text}</text>
             <path d="{path_d}" fill="none" stroke="{col_line}" stroke-width="{w_line}" />
             <circle cx="{ex}" cy="{ey}" r="{r_dot}" fill="{col_dot}" stroke="white" stroke-width="1" />
         </g>
@@ -222,17 +256,14 @@ def generate_scientific_svg(target_class):
 
     final_svg = f"""<svg id="{svg_id}" width="100%" height="100%" viewBox="0 0 800 420" xmlns="http://www.w3.org/2000/svg">
         <defs>
-            <style>
-                @import url('https://fonts.googleapis.com/css2?family=Lato:wght@400;700&display=swap');
-                text {{ font-family: 'Lato', 'Helvetica Neue', Arial, sans-serif; }}
-            </style>
+            <style>@import url('https://fonts.googleapis.com/css2?family=Lato:wght@400;700&display=swap'); text {{ font-family: 'Lato', sans-serif; }}</style>
             <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>
         <rect width="800" height="420" fill="white" />
         {shapes}
         {sec_svg}
         {label_svg}
-        <text x="400" y="400" text-anchor="middle" font-family="'Lato', 'Helvetica Neue', Arial, sans-serif" font-size="16" fill="#546E7A" font-weight="bold">Prediction: {target_class}</text>
+        <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}
@@ -252,25 +283,71 @@ def generate_scientific_svg(target_class):
     return html
 
 # ==========================
-# 4. Attention Heatmap
+# 4. Panel D: Wrapped Attention Heatmap (折行显示氨基酸)
 # ==========================
-def draw_attention_heatmap_strip(weights, sequence):
+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())
     
-    # 字体设置 (Matplotlib)
-    plt.rcParams['font.family'] = 'sans-serif'
-    plt.rcParams['font.sans-serif'] = ['Lato', 'Helvetica Neue', 'Arial']
+    seq_len = len(sequence)
+    # 计算行数
+    num_rows = (seq_len + chars_per_line - 1) // chars_per_line
     
-    fig, ax = plt.subplots(figsize=(8, 1.5), dpi=150)
-    data = weights.reshape(1, -1)
-    im = ax.imshow(data, cmap='Reds', aspect='auto', vmin=0, vmax=1)
+    # 动态调整画布高度
+    fig_height = max(2, num_rows * 0.8) 
+    fig, axes = plt.subplots(num_rows, 1, figsize=(10, fig_height), dpi=150)
     
-    ax.set_title("Sequence Attention Heatmap (Darker = Higher Attention)", fontsize=10, fontweight='bold', color='#37474F', pad=8)
-    ax.set_xlabel("Residue Position", fontsize=9)
-    ax.set_yticks([])
-    for s in ax.spines.values(): s.set_visible(False)
+    # 如果只有一行，axes不是列表，强制转列表
+    if num_rows == 1:
+        axes = [axes]
+        
+    # 字体设置
+    plt.rcParams['font.family'] = 'sans-serif'
+    plt.rcParams['font.sans-serif'] = ['Lato', 'monospace']
+
+    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...)
+        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)
     return fig
 
 # ==========================
@@ -287,38 +364,45 @@ 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]
     
-    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)}
+    # 获取原始预测 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)
     
-    svg = generate_scientific_svg(top_label)
-    heatmap = draw_attention_heatmap_strip(pooling_weights[0].cpu().numpy(), seq)
+    # 构建置信度字典 (全部转换为标准名称)
+    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 Style Fonts)
+# 6. UI Layout (UniProt Font Style)
 # ==========================
 layout_css = """
 @import url('https://fonts.googleapis.com/css2?family=Lato:wght@300;400;700&display=swap');
-
-/* 全局字体设置为 Lato */
-body, button, input, textarea, .gradio-container { 
-    font-family: 'Lato', 'Helvetica Neue', Arial, sans-serif !important; 
-}
+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-title { font-size: 2.2rem; font-weight: 800; color: #0288D1; margin-bottom: 5px; letter-spacing: -0.5px; }
+.header-title { font-size: 2.2rem; font-weight: 800; color: #0288D1; margin-bottom: 5px; }
 .header-sub { font-size: 1.0rem; color: #0277BD; }
 .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; }
 """
 
-# 设置 Gradio 主题字体为 Lato
-theme = gr.themes.Soft(
-    primary_hue="sky",
-    font=[gr.themes.GoogleFont("Lato"), "Helvetica Neue", "Arial", "sans-serif"]
-).set(body_background_fill="white", block_background_fill="white", block_border_width="0px")
+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>""")
@@ -342,7 +426,7 @@ with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok") as app:
             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>Attention Heatmap</div>")
+            gr.Markdown("<div class='panel-header'><span class='panel-label'>D</span>Attention Heatmap (Sequence Weights)</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])