Update app.py

app.py

@@ -85,71 +85,103 @@ classifier.eval()
 print("✅ Ready.")
 
 # ==========================
-# 3. Panel B: SVG 绘图引擎 (6标签标准版)
+# 3. 标签映射与 SVG 引擎 (UniProt Font)
 # ==========================
+
+def map_model_label_to_svg_key(model_label):
+    l = model_label.strip()
+    mapping = {
+        "Extracellular": "sec",
+        "OuterMembrane": "om",
+        "Periplasmic": "peri",
+        "Cellwall": "cw",
+        "CYtoplasmicMembrane": "im",
+        "Cytoplasmic": "cyto"
+    }
+    if l in mapping: return mapping[l]
+    l_lower = l.lower()
+    for k, v in mapping.items():
+        if k.lower() == l_lower: return v
+    return None
+
+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"
+    return "negative"
+
 def generate_scientific_svg(target_class):
-    target = target_class.lower() if target_class else ""
+    active_key = map_model_label_to_svg_key(target_class)
+    gram_type = infer_gram_type(target_class)
     
-    # 状态判断 (6类)
-    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_cyto = "cytoplasm" in target or "cytosol" in target
-
-    # 颜色配置 (Nature Style)
+    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")
+
     c = {
         'hl_stroke': '#D32F2F', 'hl_fill': '#FFEBEE', 'hl_text': '#B71C1C', 'hl_dot': '#D32F2F',
-        'bg_stroke': '#90A4AE', 'bg_fill': '#F9FAFB', # 极淡灰
+        'bg_stroke': '#90A4AE', 'bg_fill': '#FAFAFA', 
         'bg_text': '#78909C', 'bg_line': '#CFD8DC', 'bg_dot': '#B0BEC5'
     }
 
-    # 几何参数
     svg_id = f"svg_{str(uuid.uuid4())[:8]}"
-    cx, cy = 300, 210  # 细菌中心
-    tx = 620           # 标签 X 坐标
+    cx, cy = 300, 210
+    tx = 620
 
-    # --- 1. 绘制细菌主体 (默认使用 G- 结构以展示全要素) ---
+    # 绘制主体形状
     shapes = ""
-    
-    # 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 (Dashed)
-    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}" />'
-    
-    # Inner Membrane & Cytoplasm
-    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"
-    shapes += f'<rect x="{cx-140}" y="{cy-60}" width="280" height="120" rx="60" ry="60" fill="{fill_im}" stroke="{col_im}" stroke-width="{w_im}" />'
+    if gram_type == 'negative':
+        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}" />'
+        
+        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}" />'
+        
+        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"
+        shapes += f'<rect x="{cx-140}" y="{cy-60}" width="280" height="120" rx="60" ry="60" fill="{fill_im}" stroke="{col_im}" stroke-width="{w_im}" />'
+        
+        anchors = {
+            "sec": (cx, cy-160), "om": (cx+200, cy-60), "peri": (cx+180, cy-30),
+            "cw": (cx+170, cy), "im": (cx+140, cy+30), "cyto": (cx, cy)
+        }
+    else:
+        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"
+        shapes += f'<rect x="{cx-180}" y="{cy-100}" width="360" height="200" rx="100" ry="100" fill="{fill_bg}" stroke="{col_cw}" stroke-width="{w_cw}" stroke-opacity="0.7" />'
+        
+        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"
+        shapes += f'<rect x="{cx-140}" y="{cy-60}" width="280" height="120" rx="60" ry="60" fill="{fill_im}" stroke="{col_im}" stroke-width="{w_im}" />'
+        
+        anchors = {
+            "sec": (cx, cy-140), "om": (cx, cy),
+            "peri": (cx+160, cy-40), "cw": (cx+180, cy-60),
+            "im": (cx+140, cy+30), "cyto": (cx, cy)
+        }
 
-    # DNA Decoration
     shapes += f"""<g opacity="0.4">
         <path d="M {cx-30} {cy-10} Q {cx} {cy-50} {cx+30} {cy-10} T {cx+60} {cy}" fill="none" stroke="#CFD8DC" stroke-width="3" />
         <circle cx="{cx-40}" cy="{cy+20}" r="3" fill="#B0BEC5" /> <circle cx="{cx+20}" cy="{cy+30}" r="3" fill="#B0BEC5" />
     </g>"""
 
-    # --- 2. 标签系统 (6个完整标签 + 贝塞尔曲线) ---
-    
-    # 锚点目标坐标 (Target Anchor Points)
-    anchors = {
-        "sec":  (cx,       cy - 160), # 胞外 (悬浮)
-        "om":   (cx + 200, cy - 60),  # 外膜边界
-        "peri": (cx + 180, cy - 30),  # 周质间隙
-        "cw":   (cx + 170, cy),       # 细胞壁
-        "im":   (cx + 140, cy + 30),  # 内膜边界
-        "cyto": (cx,       cy)        # 胞质中心
-    }
+    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>
+            <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),
@@ -158,16 +190,17 @@ def generate_scientific_svg(target_class):
         ("Inner Membrane", "im", is_im),
         ("Cytoplasm", "cyto", is_cyto)
     ]
+    if gram_type == 'positive':
+        labels_config = [l for l in labels_config if l[1] != 'om']
 
     label_svg = ""
     y_start = 50
-    y_step = 60 # 间距
+    y_step = 60
     
     for i, (text, key, active) in enumerate(labels_config):
         ty = y_start + i * y_step
         ex, ey = anchors.get(key, (0,0))
         
-        # 样式
         col_txt = c['hl_text'] if active else c['bg_text']
         w_txt = "bold" if active else "normal"
         col_line = c['hl_stroke'] if active else c['bg_line']
@@ -175,31 +208,36 @@ def generate_scientific_svg(target_class):
         col_dot = c['hl_dot'] if active else c['bg_dot']
         r_dot = "5" if active else "3"
         
-        # 贝塞尔 S 形曲线
-        # c1: 从文字左侧水平延伸; c2: 向锚点垂直延伸
         c1x, c1y = tx - 80, ty
         c2x, c2y = ex + 60, ey
         path_d = f"M {tx-10} {ty-5} C {c1x} {c1y}, {c2x} {c2y}, {ex} {ey}"
         
         label_svg += f"""
         <g>
-            <text x="{tx}" y="{ty}" fill="{col_txt}" font-weight="{w_txt}" font-size="14" font-family="Arial">{text}</text>
+            <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>
             <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>
         """
 
     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>
+            <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="Arial" font-size="16" fill="#546E7A" font-weight="bold">Prediction: {target_class}</text>
+        <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>
     </svg>"""
     
-    # 嵌入 JS 下载
     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;">Download SVG</button>
+        <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) {{
@@ -214,26 +252,24 @@ def generate_scientific_svg(target_class):
     return html
 
 # ==========================
-# 4. Panel D: Attention Heatmap (纯净热图)
+# 4. Attention Heatmap
 # ==========================
 def draw_attention_heatmap_strip(weights, sequence):
-    # 归一化
     if weights.max() > 0:
         weights = (weights - weights.min()) / (weights.max() - weights.min())
     
-    fig, ax = plt.subplots(figsize=(8, 2), dpi=150) # 稍微加宽
-    data = weights.reshape(1, -1)
+    # 字体设置 (Matplotlib)
+    plt.rcParams['font.family'] = 'sans-serif'
+    plt.rcParams['font.sans-serif'] = ['Lato', 'Helvetica Neue', 'Arial']
     
-    # 绘制热图 (Reds)
+    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)
     
-    ax.set_title("Sequence Attention Heatmap (Darker = Higher Attention)", fontsize=10, fontweight='bold', color='#37474F', pad=10)
+    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([]) # 不显示 Y 轴
-    
-    # 隐藏四周边框
-    for spine in ax.spines.values(): spine.set_visible(False)
-    
+    ax.set_yticks([])
+    for s in ax.spines.values(): s.set_visible(False)
     plt.tight_layout()
     return fig
 
@@ -248,46 +284,45 @@ def predict(sequence_input):
     with torch.no_grad():
         inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=1024).to(DEVICE)
         outputs = plm_model(**inputs)
-        
-        logits, pooling_weights = classifier(
-            outputs.last_hidden_state[:, 0, :], 
-            outputs.last_hidden_state[:, 1:-1, :], 
-            inputs['attention_mask'][:, 1:-1]
-        )
+        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)}
     
-    # Panel B: SVG
     svg = generate_scientific_svg(top_label)
-    
-    # Panel D: Heatmap (纯净版)
     heatmap = draw_attention_heatmap_strip(pooling_weights[0].cpu().numpy(), seq)
     
     return confidences, svg, heatmap
 
 # ==========================
-# 6. UI Layout (4-Block Paper Style)
+# 6. UI Layout (UniProt Style Fonts)
 # ==========================
 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; }
+@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; 
+}
+
 .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-title { font-size: 2.2rem; font-weight: 800; color: #0288D1; margin-bottom: 5px; letter-spacing: -0.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; }
 """
 
-theme = gr.themes.Soft(primary_hue="sky").set(body_background_fill="white", block_background_fill="white", block_border_width="0px")
+# 设置 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")
 
 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>""")
 
-    # Row 1: A & B
     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>")
@@ -301,7 +336,6 @@ 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: C & D
     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>")