Update app.py

app.py

@@ -1,13 +1,7 @@
-# app_locpred_prok.py
 import os
-import re
 import json
+import re
 import uuid
-import io
-import shutil
-import base64
-from typing import Tuple
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -16,32 +10,30 @@ import matplotlib.pyplot as plt
 import numpy as np
 from transformers import AutoTokenizer, AutoModel
 
-# Optional server-side PDF export dependency
-try:
-    import cairosvg
-    CAIROSVG_AVAILABLE = True
-except Exception:
-    CAIROSVG_AVAILABLE = False
-
-# ========== Environment (same cache handling as before) ==========
-plt.switch_backend('Agg')
+# ==========================
+# 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"
 
+import shutil
 for path in ["/tmp/hf_cache", os.path.expanduser("~/.cache/huggingface")]:
     shutil.rmtree(path, ignore_errors=True)
     os.makedirs(path, exist_ok=True)
 
-# ========== Model architecture (same as you had) ==========
+# ==========================
+# 1. 模型架构
+# ==========================
 class AttentionPooling(nn.Module):
     def __init__(self, d_model):
         super().__init__()
         self.attention_net = nn.Linear(d_model, 1)
 
     def forward(self, x, mask):
-        attn_logits = self.attention_net(x).squeeze(2)
-        attn_logits = attn_logits.masked_fill(mask == 0, -1e9)
+        attn_logits = self.attention_net(x).squeeze(2) 
+        attn_logits.masked_fill_(mask == 0, -float('inf'))
         attn_weights = F.softmax(attn_logits, dim=1)
         return torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1), attn_weights
 
@@ -54,11 +46,7 @@ class ProtDualBranchEnhancedClassifier(nn.Module):
         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):
         z_cls = self.cls_projector(cls_embedding)
@@ -71,527 +59,259 @@ class ProtDualBranchEnhancedClassifier(nn.Module):
         z_fused_gated = z_fused_concat * gate_values
         return self.classifier_head(z_fused_gated), pooling_weights
 
-# ========== Load models (keep same variable names so your logic works) ==========
+# ==========================
+# 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 you want to test UI without model files, set MOCK_MODE = True
-MOCK_MODE = False
-
-if not MOCK_MODE:
-    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}")
-
-    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()}
-    NUM_CLASSES = len(idx_to_label)
-    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("✅ Models loaded.")
-else:
-    # Mock objects for UI development
-    idx_to_label = {0: "Cytoplasm", 1: "Inner Membrane", 2: "Periplasm", 3: "Outer Membrane", 4: "Cell Wall", 5: "Extracellular"}
-    NUM_CLASSES = len(idx_to_label)
-    tokenizer = None
-    plm_model = None
-    classifier = None
-
-# ========== SVG generator (UniProt-like) ==========
-def generate_uniprot_style_svg(pred_label: str,
-                               gram: str = "negative",
-                               theme: str = "uniprot-blue",
-                               layout: str = "circular",
-                               high_res: bool = False,
-                               uid: str = None) -> str:
-    """
-    Create a UniProt-like bacterial localization diagram:
-      - pred_label: predicted top class name (used to highlight)
-      - gram: 'negative' or 'positive'
-      - theme: 'uniprot-blue', 'red-highlight', 'auto' (auto uses prefers-color-scheme)
-      - layout: 'circular' (capsule) or 'horizontal' (cell left, labels right)
-      - high_res: True -> larger viewBox for higher-quality PNG/PDF
-    Returns: HTML string containing responsive SVG and download buttons.
-    """
-
-    target = pred_label.lower() if pred_label else ""
-    is_active = {
-        "sec":  ("extracellular" in target) or ("secreted" in target),
-        "om":   ("outer membrane" in target),
-        "peri": ("periplasm" in target),
-        "cw":   ("cell wall" in target),
-        "im":   ("inner membrane" in target) or ("plasma membrane" in target),
-        "cyto": ("cytoplasm" in target) or ("cytosol" in target)
-    }
-
-    # If gram-positive, there is no outer membrane and cell wall is thicker
-    have_outer_membrane = (gram == "negative")
-
-    # color themes using CSS variables (supports prefers-color-scheme)
-    css_vars = {
-        "uniprot-blue": {
-            "--om-fill": "#F5F7FA", "--im-fill": "#FFFFFF", "--stroke": "#607D8B",
-            "--muted": "#B0BEC5", "--text": "#263238", "--highlight": "#0288D1"
-        },
-        "red-highlight": {
-            "--om-fill": "#FFEBEE", "--im-fill": "#FFFFFF", "--stroke": "#607D8B",
-            "--muted": "#B0BEC5", "--text": "#263238", "--highlight": "#D32F2F"
-        }
-    }
-
-    selected = css_vars.get(theme if theme in css_vars else "uniprot-blue")
-
-    # Unique IDs for DOM elements to allow multiple diagrams on page
-    uid = uid or str(uuid.uuid4()).replace("-", "")[:10]
-    svg_id = f"loc_svg_{uid}"
-
-    # Sizes
-    if high_res:
-        W, H = 1600, 800
-    else:
-        W, H = 800, 420
-
-    # anchor coords and label positions (tuned for viewBox)
-    center_x = int(W * 0.38)
-    center_y = int(H * 0.5)
-    label_x = int(W * 0.76)
-    label_x_left = int(W * 0.58)
-
-    label_y_map = {
-        "sec": int(H * 0.12),
-        "om": int(H * 0.22),
-        "peri": int(H * 0.32),
-        "cw": int(H * 0.42),
-        "im": int(H * 0.62),
-        "cyto": int(H * 0.78)
+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}")
+
+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()}
+NUM_CLASSES = len(idx_to_label)
+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 绘图引擎 (6标签标准版)
+# ==========================
+def generate_scientific_svg(target_class):
+    target = target_class.lower() if target_class else ""
+    
+    # 状态判断 (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)
+    c = {
+        'hl_stroke': '#D32F2F', 'hl_fill': '#FFEBEE', 'hl_text': '#B71C1C', 'hl_dot': '#D32F2F',
+        'bg_stroke': '#90A4AE', 'bg_fill': '#F9FAFB', # 极淡灰
+        'bg_text': '#78909C', 'bg_line': '#CFD8DC', 'bg_dot': '#B0BEC5'
     }
 
+    # 几何参数
+    svg_id = f"svg_{str(uuid.uuid4())[:8]}"
+    cx, cy = 300, 210  # 细菌中心
+    tx = 620           # 标签 X 坐标
+
+    # --- 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}" />'
+
+    # 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": (center_x + 160, center_y - 160),
-        "om":  (center_x + 160, center_y - 100),
-        "peri":(center_x + 140, center_y - 40),
-        "cw":  (center_x + 120, center_y + 10),
-        "im":  (center_x + 80,  center_y + 60),
-        "cyto":(center_x + 20,  center_y + 70)
+        "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)        # 胞质中心
     }
 
-    # Build helper for connector group
-    def connector_svg(key, text):
-        ex, ey = anchors[key]
-        tx, ty = label_x, label_y_map[key]
-        # styling depends on active
-        active = is_active.get(key, False)
-        stroke_color = selected["--highlight"] if active else selected["--muted"]
-        stroke_w = "2.6" if active else "1.4"
-        fontw = "700" if active else "400"
-        dot_r = "6" if active else "4"
-        path = f"M {tx-20} {ty-6} C {tx-100} {ty-6}, {ex+60} {ey+10}, {ex} {ey}"
-        return f"""
-        <g class="connector connector-{key}">
-            <text x="{tx}" y="{ty}" fill="{selected['--text']}" font-weight="{fontw}" font-size="{14 if not high_res else 22}" font-family="Inter, Arial">{text}</text>
-            <path d="{path}" fill="none" stroke="{stroke_color}" stroke-width="{stroke_w}" stroke-linecap="round" stroke-linejoin="round" />
-            <circle cx="{ex}" cy="{ey}" r="{dot_r}" fill="{stroke_color}" stroke="white" stroke-width="{1 if not high_res else 1.6}" />
+    # 标签配置
+    labels_config = [
+        ("Extracellular", "sec", is_sec),
+        ("Outer Membrane", "om", is_om),
+        ("Periplasm", "peri", is_peri),
+        ("Cell Wall", "cw", is_cw),
+        ("Inner Membrane", "im", is_im),
+        ("Cytoplasm", "cyto", is_cyto)
+    ]
+
+    label_svg = ""
+    y_start = 50
+    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']
+        w_line = "2.5" if active else "1.0"
+        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>
+            <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>
         """
 
-    # Build cell shapes: capsule-like curves - simpler parametric shapes
-    # Outer membrane (or single outer layer for gram-positive), cell wall, inner membrane
-    # Different shapes when gram-positive: thicker cell wall, no outer membrane ring.
-    # We'll draw with bezier-like path strings tuned to look UniProt-ish.
-    if have_outer_membrane:
-        om_fill = "var(--om-fill)"
-        om_stroke = "var(--stroke)"
-        cw_fill = "none"
-        cw_stroke = "var(--muted)"
-        im_fill = "var(--im-fill)"
-        im_stroke = "var(--stroke)"
-    else:
-        # gram-positive: cell wall thicker and outer membrane absent
-        om_fill = "none"
-        om_stroke = "none"
-        cw_fill = "var(--om-fill)"
-        cw_stroke = "var(--stroke)"
-        im_fill = "var(--im-fill)"
-        im_stroke = "var(--stroke)"
-
-    # layer highlight override if active
-    def stroke_override(key, base):
-        if is_active.get(key, False):
-            return selected["--highlight"]
-        return base
-
-    # inline CSS for animations and hover effects
-    svg_style = f"""
-    <style>
-      /* theme vars */
-      :root {{
-        --om-fill: {selected['--om-fill']};
-        --im-fill: {selected['--im-fill']};
-        --stroke: {selected['--stroke']};
-        --muted: {selected['--muted']};
-        --text: {selected['--text']};
-        --highlight: {selected['--highlight']};
-      }}
-      @media (prefers-color-scheme: dark) {{
-        :root {{
-          --om-fill: #28343a;
-          --im-fill: #1f2b30;
-          --stroke: #90a4ae;
-          --muted: #546e7a;
-          --text: #e0f2f1;
-        }}
-      }}
-
-      /* connector hover: slightly thicken the path and enlarge dot */
-      .connector path {{ transition: stroke-width 180ms ease, stroke 180ms ease; opacity:0.95; }}
-      .connector circle {{ transition: r 160ms ease, transform 160ms ease; transform-origin: center; }}
-      .connector text {{ transition: fill 160ms ease; }}
-
-      /* on hover of group, emphasize */
-      .connector:hover path {{ stroke-width: calc(var(--hover-w, 3)); opacity:1; filter: drop-shadow(0 2px 2px rgba(0,0,0,0.06)); }}
-      .connector:hover circle {{ transform: scale(1.25); }}
-      /* subtle floating animation for lines */
-      .connector path {{}}
-      @keyframes floatx {{ 0% {{ transform: translateX(0px); }} 50% {{ transform: translateX(1px); }} 100% {{ transform: translateX(0px); }} }}
-      .connector path {{ animation: floatx 4s ease-in-out infinite; animation-delay: calc(var(--i, 0) * 0.12s); opacity:0.95; }}
-
-      /* make the whole svg responsive */
-      svg {{ max-width: 100%; height: auto; display:block; }}
-
-      /* layer highlight when active: add glow */
-      .layer-active {{ filter: drop-shadow(0 4px 8px rgba(0,0,0,0.08)); }}
-    </style>
-    """
-
-    # Compose SVG core shapes (simplified, but tuned coordinates)
-    # We use path shapes with translated center for convenience.
-    cell_shapes = ""
-    # Outer membrane / envelope
-    if have_outer_membrane:
-        cell_shapes += f'''
-        <g id="outer_membrane" class="layer {'layer-active' if is_active['om'] else ''}">
-          <ellipse cx="{center_x}" cy="{center_y}" rx="{220 if not high_res else 440}" ry="{170 if not high_res else 340}"
-                   fill="var(--om-fill)" stroke="{stroke_override('om', 'var(--stroke)')}" stroke-width="{3 if is_active['om'] else 2}"/>
-        </g>
-        '''
-        # Cell wall (dashed)
-        cell_shapes += f'''
-        <g id="cell_wall">
-          <ellipse cx="{center_x}" cy="{center_y}" rx="{190 if not high_res else 380}" ry="{150 if not high_res else 300}"
-                   fill="none" stroke="{stroke_override('cw','var(--muted)')}" stroke-width="{4 if is_active['cw'] else 2}" stroke-dasharray="10 6"/>
-        </g>
-        '''
-        # inner membrane
-        cell_shapes += f'''
-        <g id="inner_membrane" class="layer {'layer-active' if is_active['im'] else ''}">
-          <ellipse cx="{center_x}" cy="{center_y}" rx="{140 if not high_res else 280}" ry="{100 if not high_res else 200}"
-                   fill="var(--im-fill)" stroke="{stroke_override('im','var(--stroke)')}" stroke-width="{3 if is_active['im'] else 1.8}"/>
-        </g>
-        '''
-    else:
-        # Gram positive: thick cell wall as filled ellipse + inner membrane
-        cell_shapes += f'''
-        <g id="cell_wall_gp" class="layer {'layer-active' if is_active['cw'] else ''}">
-          <ellipse cx="{center_x}" cy="{center_y}" rx="{230 if not high_res else 460}" ry="{180 if not high_res else 360}"
-                   fill="{selected['--om-fill']}" stroke="{stroke_override('cw','var(--stroke)')}" stroke-width="{3 if is_active['cw'] else 2}"/>
-        </g>
-        <g id="inner_membrane" class="layer {'layer-active' if is_active['im'] else ''}">
-          <ellipse cx="{center_x}" cy="{center_y}" rx="{150 if not high_res else 300}" ry="{110 if not high_res else 220}"
-                   fill="var(--im-fill)" stroke="{stroke_override('im','var(--stroke)')}" stroke-width="{2 if is_active['im'] else 1.4}"/>
-        </g>
-        '''
-
-    # cytoplasm ornament
-    cell_shapes += f'''
-    <g id="cytoplasm_wiggles" opacity="0.65">
-      <path d="M {center_x-60} {center_y+10} q 30 -50 70 0 q 30 50 70 0" stroke="var(--muted)" stroke-width="6" fill="none" stroke-linecap="round"/>
-      <circle cx="{center_x-40}" cy="{center_y+40}" r="{3 if not high_res else 6}" fill="var(--muted)"/>
-      <circle cx="{center_x+20}" cy="{center_y+50}" r="{3 if not high_res else 6}" fill="var(--muted)"/>
-    </g>
-    '''
-
-    # connectors
-    connectors = ""
-    connectors += connector_svg("sec", "Extracellular / Secreted")
-    # outer membrane only if present
-    if have_outer_membrane:
-        connectors += connector_svg("om", "Outer Membrane")
-    connectors += connector_svg("peri", "Periplasm")
-    connectors += connector_svg("cw", "Cell Wall")
-    connectors += connector_svg("im", "Inner Membrane")
-    connectors += connector_svg("cyto", "Cytoplasm")
-
-    # Build download buttons and client-side JS to download SVG and PNG
-    # PDF export will call a Gradio server endpoint (provided below)
-    html = f"""
-    <div style="width:100%; text-align:center;">
-    {svg_style}
-    <svg id="{svg_id}" viewBox="0 0 {W} {H}" xmlns="http://www.w3.org/2000/svg" role="img" aria-label="Bacterial localization diagram">
-        <defs>
-            <style><![CDATA[
-                text {{ font-family: Inter, Arial, sans-serif; }}
-            ]]></style>
-        </defs>
-        {cell_shapes}
-        {connectors}
-    </svg>
-
-    <div style="margin-top:8px; display:flex; gap:8px; justify-content:center; align-items:center;">
-      <button id="download_svg_{uid}" class="download-btn">Download SVG</button>
-      <button id="download_png_{uid}" class="download-btn">Download PNG</button>
-      <button id="download_pdf_{uid}" class="download-btn">Download PDF</button>
-      <div style="font-size:12px; color:var(--text); align-self:center;">{gram.title()} · {layout.title()} {'· High-res' if high_res else ''}</div>
-    </div>
+    final_svg = f"""<svg id="{svg_id}" width="100%" height="100%" viewBox="0 0 800 420" xmlns="http://www.w3.org/2000/svg">
+        <rect width="800" height="420" fill="white" />
+        {shapes}
+        {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>
+    </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>
     </div>
-
     <script>
-    (function(){{
-      const svgEl = document.getElementById("{svg_id}");
-      const btnSvg = document.getElementById("download_svg_{uid}");
-      const btnPng = document.getElementById("download_png_{uid}");
-      const btnPdf = document.getElementById("download_pdf_{uid}");
-
-      function downloadFile(filename, blob) {{
-        const url = URL.createObjectURL(blob);
-        const a = document.createElement('a');
-        a.href = url; a.download = filename; document.body.appendChild(a); a.click();
-        setTimeout(()=>{{ URL.revokeObjectURL(url); a.remove(); }}, 200);
-      }}
-
-      btnSvg.addEventListener('click', ()=>{{
-        const serializer = new XMLSerializer();
-        let source = serializer.serializeToString(svgEl);
-        if(!source.match(/^<svg[^>]+xmlns="http:\\/\\/www.w3.org\\/2000\\/svg"/)) {{
-            source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
-        }}
-        const blob = new Blob([source], {{type: 'image/svg+xml;charset=utf-8'}});
-        downloadFile('locpred_diagram.svg', blob);
-      }});
-
-      btnPng.addEventListener('click', ()=>{{
-        const serializer = new XMLSerializer();
-        let source = serializer.serializeToString(svgEl);
-        if(!source.match(/^<svg[^>]+xmlns="http:\\/\\/www.w3.org\\/2000\\/svg"/)) {{
-            source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
-        }}
-        const svgBlob = new Blob([source], {{type: 'image/svg+xml;charset=utf-8'}});
-        const url = URL.createObjectURL(svgBlob);
-        const img = new Image();
-        img.onload = function() {{
-          const canvas = document.createElement('canvas');
-          // scale 2x for higher quality
-          const scale = 2;
-          canvas.width = img.width * scale;
-          canvas.height = img.height * scale;
-          const ctx = canvas.getContext('2d');
-          // optional white background
-          ctx.fillStyle = "white";
-          ctx.fillRect(0,0,canvas.width,canvas.height);
-          ctx.drawImage(img, 0, 0, canvas.width, canvas.height);
-          canvas.toBlob(function(blob) {{
-            downloadFile('locpred_diagram.png', blob);
-          }}, 'image/png');
-          URL.revokeObjectURL(url);
-        }};
-        img.onerror = function(e) {{
-          alert('Failed to render PNG in your browser.');
-          URL.revokeObjectURL(url);
-        }};
-        img.src = url;
-      }});
-
-      btnPdf.addEventListener('click', async ()=>{{
-        // send the SVG string to the server /gradio route for PDF conversion
-        const serializer = new XMLSerializer();
-        let source = serializer.serializeToString(svgEl);
-        if(!source.match(/^<svg[^>]+xmlns="http:\\/\\/www.w3.org\\/2000\\/svg"/)) {{
-            source = source.replace(/^<svg/, '<svg xmlns="http://www.w3.org/2000/svg"');
-        }}
-        // call Gradio server function via fetch to /convert_svg_to_pdf (provided below)
-        try {{
-          const resp = await fetch('/convert_svg_to_pdf', {{
-            method: 'POST',
-            headers: {{ 'Content-Type': 'application/json' }},
-            body: JSON.stringify({{ svg: source }})
-          }});
-          if(!resp.ok) {{
-            const txt = await resp.text();
-            alert('PDF conversion failed: ' + txt);
-            return;
-          }}
-          const blob = await resp.blob();
-          downloadFile('locpred_diagram.pdf', blob);
-        }} catch (err) {{
-          alert('PDF conversion failed: ' + err);
+        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>
-    """
-
+    </script></div>"""
     return html
 
-# ========== Heatmap (same as before) ==========
+# ==========================
+# 4. Panel D: 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)
-    fig, ax = plt.subplots(figsize=(8, 1.5), dpi=150)
+    
+    # 绘制热图 (Reds)
     im = ax.imshow(data, cmap='Reds', aspect='auto', vmin=0, vmax=1)
-    ax.set_title('Sequence Attention Heatmap (High Color = Key Feature)', fontsize=10, fontweight='bold', color='#37474F', pad=6)
-    ax.set_xlabel('Residue Position', fontsize=9)
-    ax.set_yticks([])
-    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)
+    
+    ax.set_title("Sequence Attention Heatmap (Darker = Higher Attention)", fontsize=10, fontweight='bold', color='#37474F', pad=10)
+    ax.set_xlabel("Residue Position", fontsize=9)
+    ax.set_yticks([]) # 不显示 Y 轴
+    
+    # 隐藏四周边框
+    for spine in ax.spines.values(): spine.set_visible(False)
+    
     plt.tight_layout()
     return fig
 
-# ========== Prediction function ==========
-def predict(sequence_input: str, gram_choice: str, theme_choice: str, layout_choice: str, high_res_flag: bool):
-    """
-    Returns:
-      - confidences dict (for Label)
-      - svg html string (for HTML output)
-      - attention heatmap figure (for Plot)
-    """
-    if not sequence_input or sequence_input.isspace():
-        raise gr.Error("Empty Input")
-
+# ==========================
+# 5. 预测主逻辑
+# ==========================
+def predict(sequence_input):
+    if not sequence_input or sequence_input.isspace(): raise gr.Error("Empty Input")
     seq = "".join(sequence_input.split('\n')[1:]) if sequence_input.startswith('>') else sequence_input
     seq = re.sub(r'[^A-Z]', '', seq.upper())[:1024]
-    if not seq:
-        raise gr.Error("Invalid Sequence")
-
-    if MOCK_MODE:
-        # mock probabilities
-        probs = torch.softmax(torch.randn(NUM_CLASSES), dim=0)
-        logits = probs
-        pooling_weights = np.abs(np.random.randn(len(seq))).astype(float)
-        pooling_weights = pooling_weights / pooling_weights.sum()
-        top_label = idx_to_label[int(torch.argmax(probs).item())]
-    else:
-        with torch.no_grad():
-            inputs = tokenizer(seq, return_tensors="pt", truncation=True, max_length=1024).to(DEVICE)
-            outputs = plm_model(**inputs)
-            hidden_states = outputs.last_hidden_state
-            cls_embedding = hidden_states[:, 0, :]
-            token_embeddings = hidden_states[:, 1:-1, :]
-            token_mask = inputs['attention_mask'][:, 1:-1]
-            logits, pooling_weights = classifier(cls_embedding, token_embeddings, token_mask)
-            probs = F.softmax(logits, dim=1)[0]
-            top_label = idx_to_label[torch.argmax(probs).item()]
-            pooling_weights = pooling_weights[0].cpu().numpy()
-
-    confidences = { idx_to_label[i]: float(p) for i,p in enumerate(probs) } if not MOCK_MODE else { idx_to_label[i]: float(p) for i,p in enumerate(probs)}
-    svg_html = generate_uniprot_style_svg(top_label, gram=gram_choice, theme=theme_choice, layout=layout_choice, high_res=high_res_flag)
-    heatmap_fig = draw_attention_heatmap_strip(np.array(pooling_weights), seq)
-
-    return confidences, svg_html, heatmap_fig
-
-# ========== Server-side PDF conversion endpoint for Gradio ==========
-# This function will be exposed at /convert_svg_to_pdf when app launches.
-def convert_svg_to_pdf_endpoint(svg_str: str):
-    """
-    Convert an SVG string to PDF bytes using cairosvg (if available).
-    Return bytes-like object (PDF) or raise error.
-    """
-    if not CAIROSVG_AVAILABLE:
-        raise RuntimeError("Server-side PDF conversion requires 'cairosvg' package. Install with: pip install cairosvg")
-
-    # cairosvg.svg2pdf can take bytes or string
-    pdf_bytes = cairosvg.svg2pdf(bytestring=svg_str.encode('utf-8'))
-    return ("locpred_diagram.pdf", pdf_bytes)
-
-# ========== UI layout (Gradio) ==========
+    
+    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]
+        )
+        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)
+# ==========================
 layout_css = """
-/* small customizations and CSS variables fallback */
-:root { --panel-bg: white; }
-.gradio-container { max-width: 1200px; margin: 0 auto; }
-.download-btn { padding:8px 12px; border-radius:6px; border:1px solid #D1D5DB; background:transparent; cursor:pointer; }
-.download-btn:hover { box-shadow: 0 4px 14px rgba(16,24,40,0.08); }
-.panel-card { border:1px solid #e6eef5; border-radius:8px; padding:12px; background:var(--panel-bg); }
-.panel-header { font-weight:700; color:#475569; border-bottom:2px solid #f1f5f9; padding-bottom:8px; margin-bottom:10px; }
+@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-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; }
+.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")
 
-with gr.Blocks(theme=theme, css=layout_css, title="LocPred-Prok (UniProt-style)") as app:
-    gr.Markdown("<div style='font-size:22px; font-weight:800; color:#0288D1;'>LocPred-Prok — UniProt-style visualization</div>")
+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(scale=6):
-            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>A</span>Sequence Input</div>")
-            sequence_input = gr.Textbox(lines=8, placeholder=">Sequence (single-letter amino acids) or paste raw sequence", show_label=False)
+        with gr.Column(elem_classes="panel-card"):
+            gr.Markdown("<div class='panel-header'><span class='panel-label'>A</span>Sequence Input</div>")
+            sequence_input = gr.Textbox(lines=8, show_label=False, placeholder=">Sequence...")
             with gr.Row():
                 clear_btn = gr.ClearButton(sequence_input, value="Clear")
                 submit_btn = gr.Button("Predict Analysis", variant="primary")
-            with gr.Row():
-                gram_choice = gr.Radio(choices=["negative", "positive"], value="negative", label="Gram type")
-                theme_choice = gr.Radio(choices=["uniprot-blue", "red-highlight"], value="uniprot-blue", label="Color Theme")
-                layout_choice = gr.Radio(choices=["circular", "horizontal"], value="circular", label="Diagram Layout")
-                high_res_flag = gr.Checkbox(value=False, label="High resolution (bigger SVG/PDF)")
-            gr.Examples([[">Outer Membrane\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAF..."]], inputs=sequence_input)
+            gr.Examples([[">Outer Membrane\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAFGGYQVNPYVGFEMGYDWLGRMPYKGSVENGAYKAQGVQLTAKLGYPITDDLDIYTRLGGMVWRADTKSNVYGKNHDTGVSPVFAGGVEYAITPEIATRLEYQWTNNIGDAHTIGTRPDNGMLSLGVSYRFGQGEAAPVVAPAPAPAPEVQTKHFTLKSDVLFNFNKATLKPEGQAALDQLYSQLSNLDPKDGSVVVLGYTDRIGSDAYNQGLSERRAQSVVDYLISKGIPADKISARGMGESNPVTGNTCDNVKQRAALIDCLAPDRRVEIEVKGIKDVVTQPQA"]], inputs=sequence_input, label=None)
 
-        with gr.Column(scale=6):
-            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>B</span>Localization Visualization</div>")
+        with gr.Column(elem_classes="panel-card"):
+            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(scale=6):
-            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>C</span>Prediction Confidence</div>")
+        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(scale=6):
-            gr.Markdown("<div class='panel-header'><span style='background:#E0F7FA;color:#0277BD;padding:3px 6px;border-radius:4px;font-weight:800;margin-right:8px;'>D</span>Attention Heatmap</div>")
-            output_plot = gr.Plot(show_label=False)
 
-    submit_btn.click(fn=predict, inputs=[sequence_input, gram_choice, theme_choice, layout_choice, high_res_flag], outputs=[output_label, output_svg, output_plot])
-    clear_btn.click(lambda: [None, None, None], outputs=[output_label, output_svg, output_plot])
-
-    # Expose PDF conversion endpoint: Gradio allows adding a separate route handler via app.launch later.
-    # We'll attach the endpoint to the FastAPI app used by Gradio when launching.
-
-# ========== Run server with custom route for PDF conversion ==========
-if __name__ == "__main__":
-    from fastapi import FastAPI, Request, Response
-    import uvicorn
-
-    # Build Gradio app to get underlying FastAPI instance
-    demo = app
-    # Get the underlying FastAPI app (gradio >= 3.0)
-    # When launching, we'll mount a custom route /convert_svg_to_pdf handled by convert_svg_to_pdf_endpoint
-    # Gradio's launch will create a FastAPI object; to avoid internal changes, we use the `server_name` arg.
+        with gr.Column(elem_classes="panel-card"):
+            gr.Markdown("<div class='panel-header'><span class='panel-label'>D</span>Attention Heatmap</div>")
+            output_plot = gr.Plot(label="Attention", show_label=False)
 
-    # Create a lightweight FastAPI for the PDF endpoint and mount the Gradio interface into it
-    fast_app = FastAPI()
-
-    @fast_app.post("/convert_svg_to_pdf")
-    async def convert_svg_to_pdf_api(request: Request):
-        payload = await request.json()
-        svg = payload.get("svg", None)
-        if not svg:
-            return Response(content="No svg provided", status_code=400)
-        if not CAIROSVG_AVAILABLE:
-            return Response(content="Server-side PDF conversion unavailable: install 'cairosvg' in the server environment.", status_code=501)
-        try:
-            pdf_bytes = cairosvg.svg2pdf(bytestring=svg.encode('utf-8'))
-            return Response(content=pdf_bytes, media_type="application/pdf")
-        except Exception as e:
-            return Response(content=f"PDF conversion error: {e}", status_code=500)
-
-    # Mount the Gradio interface at root
-    gr.mount_gradio_app(fast_app, demo, path="/")
+    submit_btn.click(fn=predict, inputs=sequence_input, outputs=[output_label, output_svg, output_plot])
+    clear_btn.click(lambda: [None, None, None], outputs=[output_label, output_svg, output_plot])
 
-    # Launch uvicorn with the FastAPI app
-    uvicorn.run(fast_app, host="0.0.0.0", port=7860, log_level="info")
+app.launch()