Update app.py

app.py

@@ -5,36 +5,32 @@ import torch.nn.functional as F
 import gradio as gr
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
-import matplotlib.patheffects as path_effects
+import matplotlib.path as mpath
 import numpy as np
 from io import BytesIO
 from PIL import Image
 from transformers import AutoTokenizer, AutoModel
 
 # ==========================
-# 0. 环境初始化与设置
+# 0. 环境初始化
 # ==========================
-# 强制使用非交互式后端，防止 matplotlib 在服务器上报错
-plt.switch_backend('Agg')
+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"
 
-# 清理缓存以防止空间溢出
 for path in ["/tmp/hf_cache", os.path.expanduser("~/.cache/huggingface")]:
     shutil.rmtree(path, ignore_errors=True)
     os.makedirs(path, exist_ok=True)
 
 # ==========================
-# 1. 模型架构定义
+# 1. 模型定义 (保持不变)
 # ==========================
 class AttentionPooling(nn.Module):
-    """Attention Pooling Layer"""
     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.masked_fill_(mask == 0, -float('inf'))
@@ -42,68 +38,42 @@ class AttentionPooling(nn.Module):
         return torch.bmm(attn_weights.unsqueeze(1), x).squeeze(1)
 
 class ProtDualBranchEnhancedClassifier(nn.Module):
-    """Enhanced dual-branch model architecture"""
     def __init__(self, d_model, projection_dim, num_classes, dropout, kernel_size):
         super().__init__()
-        # Branch 1: Global Semantic (CLS token)
         self.cls_projector = nn.Linear(d_model, projection_dim)
-        
-        # Branch 2: Local Structural (Token embeddings + CNN)
         self.token_refiner = nn.Sequential(
-            nn.Conv1d(d_model, d_model, kernel_size, padding='same'),
-            nn.ReLU()
+            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)
-        
-        # Fusion Gate
         fused_dim = projection_dim * 2
-        self.gate = nn.Sequential(
-            nn.Linear(fused_dim, fused_dim),
-            nn.Sigmoid()
-        )
-        
-        # Classifier Head
+        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)
+            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):
-        # Process Branch 1
         z_cls = self.cls_projector(cls_embedding)
-        
-        # Process Branch 2
         tok_emb_permuted = token_embeddings.permute(0, 2, 1)
         refined_tok_emb = self.token_refiner(tok_emb_permuted).permute(0, 2, 1)
         z_tok_pooled = self.attention_pooling(refined_tok_emb, mask)
         z_tok = self.tok_projector(z_tok_pooled)
-        
-        # Fusion
         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)
 
 # ==========================
-# 2. 加载模型与资源
+# 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"Error: Missing '{LABEL_MAP_PATH}'. Please upload it to your Space.")
-if not os.path.exists(CLASSIFIER_PATH):
-    raise FileNotFoundError(f"Error: Could not find '{CLASSIFIER_PATH}'. Please upload your trained .pth file.")
+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()}
@@ -111,42 +81,40 @@ with open(LABEL_MAP_PATH, 'r') as f:
 NUM_CLASSES = len(idx_to_label)
 D_MODEL = 640
 
-print(f"🔹 Loading ESM-2 Model ({PLM_MODEL_NAME})...")
+print("🔹 Loading models...")
 tokenizer = AutoTokenizer.from_pretrained(PLM_MODEL_NAME)
-plm_model = AutoModel.from_pretrained(PLM_MODEL_NAME).to(DEVICE)
-plm_model.eval()
-
-print("🔹 Loading Custom Classifier...")
-classifier = ProtDualBranchEnhancedClassifier(
-    d_model=D_MODEL, projection_dim=32, num_classes=NUM_CLASSES,
-    dropout=0.3, kernel_size=3
-).to(DEVICE)
-
+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("✅ All Models Loaded Successfully.")
+print("✅ Ready.")
 
 # ==========================
-# 3. 高保真绘图引擎 (Visualization Engine)
+# 3. UniProt 风格绘图引擎 (Schematic Vector Style)
 # ==========================
-def draw_prokaryotic_cell_detailed(target_class):
+def draw_uniprot_style_cell(target_class):
     """
-    Generate a high-fidelity, scientific illustration of a Gram-negative bacterium
-    with the predicted subcellular location highlighted.
+    绘制类似 UniProt/SwissBioPics 的扁平化矢量风格细胞图。
+    特点：清晰的黑色描边，柔和的填充色，双层膜结构。
     """
     target = target_class.lower() if target_class else ""
     
-    # --- 精细配色板 (Scientific Palette) ---
+    # === UniProt/SwissBioPics 配色板 ===
+    # 柔和、低饱和度背景，高饱和度高亮
     c = {
-        'bg': '#ffffff',
-        'om_out': '#475569', 'om_in': '#94a3b8', 'om_fill': '#e2e8f0',
-        'cw_mesh': '#64748b', 'cw_bg': '#f1f5f9',
-        'im_out': '#64748b', 'im_in': '#cbd5e1', 'im_fill': '#f8fafc',
-        'cyto_bg': '#ffffff', 'ribo': '#e2e8f0', 'dna': '#cbd5e1',
-        'hl_core': '#ef4444', 'hl_glow1': '#fca5a5', 'hl_glow2': '#fee2e2'
+        'stroke': '#263238',       # 深灰/黑色描边 (清晰感来源)
+        'cyto_bg': '#FFF9C4',      # 极淡的黄色 (胞质)
+        'peri_bg': '#E1F5FE',      # 极淡的蓝色 (周质)
+        'ext_bg': '#ffffff',       # 纯白背景
+        'mem_line': '#546E7A',     # 膜的线条颜色
+        'cw_pattern': '#90A4AE',   # 细胞壁颜色
+        'dna': '#B0BEC5',          # DNA 灰色
+        'ribo': '#CFD8DC',         # 核糖体 浅灰
+        'active': '#D32F2F',       # UniProt 经典红 (高亮)
+        'active_fill': '#FFCDD2'   # 红色填充
     }
 
-    # --- 匹配逻辑 ---
+    # 状态判断
     is_secreted = "extracellular" in target or "secreted" in target
     is_om = "outer membrane" in target
     is_peri = "periplasm" in target
@@ -154,117 +122,127 @@ def draw_prokaryotic_cell_detailed(target_class):
     is_im = "plasma membrane" in target or "inner membrane" in target
     is_cyto = "cytoplasm" in target or "cytosol" in target
 
-    # 创建高分辨率画布
-    fig, ax = plt.subplots(figsize=(8, 5), dpi=150)
-    ax.set_xlim(-1, 11)
-    ax.set_ylim(-0.5, 6.5)
+    # 创建画布 (宽高比 16:9)
+    fig, ax = plt.subplots(figsize=(8, 4.5), dpi=150)
+    ax.set_xlim(0, 10)
+    ax.set_ylim(0, 6)
     ax.axis('off')
 
-    # --- 辅助绘图函数：拟态3D胶囊 ---
-    def draw_3d_capsule(x, y, w, h, fill_c, edge_c_out, edge_c_in, z, is_hl=False, hl_style='rim'):
-        r = h / 2
-        # 基础填充
-        base = patches.FancyBboxPatch((x, y), w-h, h, boxstyle=f"round,pad=0,rounding_size={r}", fc=fill_c, ec='none', zorder=z)
-        ax.add_patch(base)
+    # === 绘图辅助函数 ===
+    def draw_membrane_layer(x, y, w, h, z, is_active):
+        """绘制双层膜结构 (两条平行线)"""
+        color = c['active'] if is_active else c['stroke']
+        lw = 2.5 if is_active else 1.2
         
-        # 内圈描边 (模拟厚度/立体感)
-        inner_border = patches.FancyBboxPatch((x+0.08, y+0.08), w-h-0.16, h-0.16, boxstyle=f"round,pad=0,rounding_size={r-0.08}", fc='none', ec=edge_c_in, lw=1.5, zorder=z+0.1)
-        ax.add_patch(inner_border)
-        
-        # 外圈描边
-        outer_border = patches.FancyBboxPatch((x, y), w-h, h, boxstyle=f"round,pad=0,rounding_size={r}", fc='none', ec=edge_c_out, lw=2, zorder=z+0.2)
-        ax.add_patch(outer_border)
-
-        # 高亮发光逻辑
-        if is_hl:
-            if hl_style == 'rim': # 膜结构边缘发光
-                for lw, alpha, color in [(6, 0.3, c['hl_glow2']), (4, 0.5, c['hl_glow1']), (2.5, 0.8, c['hl_core'])]:
-                    hl = patches.FancyBboxPatch((x, y), w-h, h, boxstyle=f"round,pad=0,rounding_size={r}", fc='none', ec=color, lw=lw, alpha=alpha, zorder=z+0.5)
-                    ax.add_patch(hl)
-            elif hl_style == 'fill': # 区域填充发光
-                hl_fill = patches.FancyBboxPatch((x, y), w-h, h, boxstyle=f"round,pad=0,rounding_size={r}", fc=c['hl_glow2'], ec=c['hl_core'], lw=2, alpha=0.8, zorder=z+0.5)
-                ax.add_patch(hl_fill)
-
-    # --- 1. 绘制层级结构 (从外向内) ---
-    
-    # Layer 1: Outer Membrane (外膜)
-    draw_3d_capsule(0, 0, 10, 6, c['om_fill'], c['om_out'], c['om_in'], z=1, is_hl=is_om, hl_style='rim')
-
-    # Layer 2: Cell Wall Space (细胞壁区域)
-    cw_base = patches.FancyBboxPatch((0.5, 0.5), 9-5, 5, boxstyle=f"round,pad=0,rounding_size={2.5}", fc=c['cw_bg'], ec='none', zorder=2)
-    ax.add_patch(cw_base)
+        # 外层线
+        p1 = patches.FancyBboxPatch((x, y), w-h, h, boxstyle=f"round,pad=0,rounding_size={h/2}", 
+                                    fc="none", ec=color, lw=lw, zorder=z)
+        ax.add_patch(p1)
+        # 内层线 (稍微小一圈，模拟双层膜厚度)
+        gap = 0.15
+        p2 = patches.FancyBboxPatch((x+gap, y+gap), w-h-2*gap, h-2*gap, boxstyle=f"round,pad=0,rounding_size={(h-2*gap)/2}", 
+                                    fc="none", ec=color, lw=lw, zorder=z)
+        ax.add_patch(p2)
+
+    # === 1. 背景层 (Backgrounds) ===
+
+    # 外膜区域背景 (Periplasm starts here)
+    # 整个胶囊的底色设为周质颜色 (Light Blue)
+    base_capsule = patches.FancyBboxPatch((1, 1), 8, 4, boxstyle="round,pad=0,rounding_size=2", 
+                                          fc=c['peri_bg'], ec="none", zorder=1)
+    ax.add_patch(base_capsule)
+
+    # 内膜区域背景 (Cytoplasm starts here)
+    # 内部胶囊底色设为胞质颜色 (Light Yellow)
+    cyto_capsule = patches.FancyBboxPatch((1.6, 1.6), 6.8, 2.8, boxstyle="round,pad=0,rounding_size=1.4", 
+                                          fc=c['cyto_bg'], ec="none", zorder=2)
+    ax.add_patch(cyto_capsule)
+
+    # === 2. 结构层 (Structures with Strokes) ===
+
+    # A. Outer Membrane (外膜) - 双层线
+    draw_membrane_layer(1, 1, 8+2, 4, z=5, is_active=is_om)
+
+    # B. Cell Wall (细胞壁) - 网格纹理层
+    # 在周质层中间画一层
+    cw_color = c['active'] if is_cw else c['cw_pattern']
+    cw_lw = 2 if is_cw else 0.5
+    cw_alpha = 1.0 if is_cw else 0.4
+    cw_patch = patches.FancyBboxPatch((1.3, 1.3), 7.4, 3.4, boxstyle="round,pad=0,rounding_size=1.7",
+                                      fc="none", ec=cw_color, lw=cw_lw, linestyle=(0, (5, 5)), zorder=4)
+    ax.add_patch(cw_patch)
     
-    # 细胞壁网格纹理 (Hatch)
-    cw_mesh_color = c['hl_core'] if is_cw else c['cw_mesh']
-    cw_mesh_lw = 1.5 if is_cw else 0.5
-    cw_texture = patches.FancyBboxPatch((0.5, 0.5), 9-5, 5, boxstyle=f"round,pad=0,rounding_size={2.5}", fc='none', ec=cw_mesh_color, lw=cw_mesh_lw, hatch='XXXX', alpha=0.3 if not is_cw else 0.6, zorder=2.1)
-    ax.add_patch(cw_texture)
-
-    # 如果是周质 (Periplasm)，在细胞壁区域加一层发光填充
-    if is_peri: 
-        draw_3d_capsule(0.4, 0.4, 9.2, 5.2, 'none', 'none', 'none', z=2.2, is_hl=True, hl_style='fill')
-
-    # Layer 3: Inner Membrane (内膜)
-    draw_3d_capsule(1.0, 1.0, 8, 4, c['im_fill'], c['im_out'], c['im_in'], z=3, is_hl=is_im, hl_style='rim')
-
-    # Layer 4: Cytoplasm (胞质)
-    cyto_fill_c = c['hl_glow2'] if is_cyto else c['cyto_bg']
-    cyto_ec = c['hl_core'] if is_cyto else 'none'
-    draw_3d_capsule(1.3, 1.3, 7.4, 3.4, cyto_fill_c, cyto_ec, 'none', z=4, is_hl=is_cyto, hl_style='fill')
-
-    # --- 2. 填充细节内容 ---
-
-    # Ribosomes (核糖体颗粒)
-    rng = np.random.default_rng(101)
-    x_ribo = rng.uniform(2.0, 8.0, 400)
-    y_ribo = rng.uniform(1.8, 4.2, 400)
-    # 简单的椭圆过滤，只保留胶囊内部的点
-    mask = ((x_ribo - 5)**2 / 3**2 + (y_ribo - 3)**2 / 1.2**2) < 1
-    ax.scatter(x_ribo[mask], y_ribo[mask], s=3, c=c['ribo'], alpha=0.6, zorder=4.1)
-
-    # Nucleoid DNA (拟核 DNA 缠绕线)
-    t = np.linspace(0, 20, 500)
-    x_dna = 5 + 2.0 * np.cos(t) * np.sin(t*0.3) * np.cos(t*0.1)
-    y_dna = 3 + 0.8 * np.sin(t) * np.cos(t*0.2)
-    # 使用 PathEffects 增加投影立体感
-    dna_line = ax.plot(x_dna, y_dna, color=c['dna'], alpha=0.8, linewidth=2, zorder=4.2)[0]
-    dna_line.set_path_effects([
-        path_effects.SimpleLineShadow(offset=(1, -1), shadow_color=c['im_fill'], alpha=0.5, linewidth=2.5),
-        path_effects.Normal()
-    ])
-
-    # Secreted Proteins (分泌蛋白)
+    # 如果是周质高亮，填充浅红色
+    if is_peri:
+         peri_hl = patches.FancyBboxPatch((1.1, 1.1), 7.8, 3.8, boxstyle="round,pad=0,rounding_size=1.9", 
+                                          fc=c['active_fill'], ec="none", alpha=0.5, zorder=1.5)
+         ax.add_patch(peri_hl)
+
+    # C. Inner Membrane (内膜) - 双层线
+    draw_membrane_layer(1.6, 1.6, 6.8, 2.8, z=6, is_active=is_im)
+
+    # D. Cytoplasm (胞质) - 如果高亮则填充红
+    if is_cyto:
+        cyto_hl = patches.FancyBboxPatch((1.7, 1.7), 6.6, 2.6, boxstyle="round,pad=0,rounding_size=1.3", 
+                                         fc=c['active_fill'], ec="none", zorder=2.1)
+        ax.add_patch(cyto_hl)
+
+    # === 3. 内容细节 (Details) ===
+
+    # DNA (拟核) - 清晰的线条，不再是模糊的
+    t = np.linspace(0, 15, 300)
+    x_dna = 5 + 2.0 * np.cos(t) * np.sin(t*0.5)
+    y_dna = 3 + 0.6 * np.sin(t)
+    ax.plot(x_dna, y_dna, color=c['dna'], lw=2, zorder=3, solid_capstyle='round')
+
+    # Ribosomes (核糖体) - 规则的小圆圈，而不是随机噪点
+    # 使用规律的分布让画面更整洁
+    rng = np.random.default_rng(42)
+    for _ in range(30):
+        rx, ry = rng.uniform(2.5, 7.5), rng.uniform(2.2, 3.8)
+        # 避开 DNA 区域 (简单判断)
+        if not (4 < rx < 6 and 2.5 < ry < 3.5): 
+            circle = patches.Circle((rx, ry), radius=0.06, fc=c['ribo'], ec="none", zorder=3)
+            ax.add_patch(circle)
+
+    # Secreted Proteins (胞外)
     if is_secreted:
-        for i in range(5):
-            sx, sy = rng.uniform(0, 10), rng.uniform(5.8, 6.5)
-            ax.scatter(sx, sy, s=50, c=c['hl_core'], edgecolors='white', zorder=10)
-            # 运动尾迹
-            ax.plot([sx, sx-0.5], [sy-0.1, sy-0.6], color=c['hl_core'], lw=1, alpha=0.6, zorder=9)
-        ax.text(5, 6.2, "Secreted / Extracellular", ha='center', va='bottom', color=c['hl_core'], fontweight='bold', fontsize=10)
-
-    # --- 3. 图例与标注 ---
-    legend_labels = [
-        ("Outer Membrane", is_om), 
-        ("Cell Wall (Peptidoglycan)", is_cw),
-        ("Periplasm", is_peri), 
-        ("Inner Membrane", is_im),
-        ("Cytoplasm & Nucleoid", is_cyto)
+        # 画几个 UniProt 风格的小球+箭头
+        ax.text(5, 5.5, "Secreted", ha='center', va='bottom', color=c['active'], fontweight='bold', fontsize=11)
+        for i in range(3):
+            sx = 4 + i
+            sy = 5.2
+            # 箭头
+            ax.arrow(sx, 4.8, 0, 0.4, head_width=0.15, head_length=0.15, fc=c['active'], ec=c['active'], width=0.04, zorder=10)
+            # 蛋白球
+            circle = patches.Circle((sx, 4.7), radius=0.15, fc=c['active'], ec=c['stroke'], zorder=10)
+            ax.add_patch(circle)
+
+    # === 4. 标注 (Annotations) ===
+    # 使用简单的线条指引
+    labels = [
+        ("Outer Membrane", 1, 5, is_om),
+        ("Periplasm", 1, 4.2, is_peri),
+        ("Cell Wall", 1, 3.6, is_cw),
+        ("Inner Membrane", 1, 3.0, is_im),
+        ("Cytoplasm", 1, 2.0, is_cyto)
     ]
     
-    ly = 5.5
-    for text, active in legend_labels:
-        lc = c['hl_core'] if active else c['im_out']
-        lw = 'bold' if active else 'normal'
-        ax.text(10.5, ly, text, ha='right', va='center', fontsize=9, color=lc, fontweight=lw)
+    for txt, x_anchor, y_anchor, active in labels:
         if active:
-            ax.scatter(10.7, ly, s=30, c=c['hl_core'], zorder=10, clip_on=False)
-        ly -= 0.5
-
-    # 底部标题
-    pred_text = f"Predicted Localization: {target_class}" if target_class else "Prokaryotic Cell Structure"
-    ax.text(5, -0.3, pred_text, ha='center', va='top', fontsize=11, fontweight='bold', color="#334155")
-
-    # 输出图片
+            # 激活状态：红色加粗字体 + 指示线
+            ax.text(0.2, y_anchor, txt, ha='left', va='center', color=c['active'], fontweight='bold', fontsize=10)
+            # 画一条线连到细胞对应位置
+            con_x = 1.0 # 胶囊左边缘附近
+            ax.plot([0.8, con_x], [y_anchor, 3], color=c['active'], lw=1, ls='dotted', alpha=0.5) 
+        else:
+            # 非激活状态：灰色小字
+            ax.text(0.2, y_anchor, txt, ha='left', va='center', color='#90A4AE', fontsize=8)
+
+    # 标题
+    ax.text(5, 0.2, f"Localization: {target_class}", ha='center', va='bottom', fontsize=11, fontweight='bold', color=c['stroke'])
+
+    # 保存
     buf = BytesIO()
     plt.savefig(buf, format='png', bbox_inches='tight', transparent=True, dpi=150)
     buf.seek(0)
@@ -273,204 +251,135 @@ def draw_prokaryotic_cell_detailed(target_class):
     return img
 
 # ==========================
-# 4. 预测主逻辑
+# 4. 预测逻辑
 # ==========================
 def predict(sequence_input):
     if not sequence_input or sequence_input.isspace():
         raise gr.Error("Please input a protein sequence.")
-    
-    # 清洗 FASTA 头和非氨基酸字符
     seq = "".join(sequence_input.split('\n')[1:]) if sequence_input.startswith('>') else sequence_input
-    seq = re.sub(r'[^A-Z]', '', seq.upper())
-    
-    if not seq: 
-        raise gr.Error("Invalid Amino Acid Sequence. Please check your input.")
-    
-    # 截断过长序列以防止 OOM
-    if len(seq) > 1024:
-        seq = seq[:1024]
+    seq = re.sub(r'[^A-Z]', '', seq.upper())[:1024]
+    if not seq: raise gr.Error("Invalid Amino Acid Sequence.")
     
     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 = classifier(cls_embedding, token_embeddings, token_mask)
+        logits = 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_prob, top_idx = torch.max(probs, dim=0)
     top_label = idx_to_label[top_idx.item()]
-    
     confidences = {idx_to_label[i]: float(p) for i, p in enumerate(probs)}
     
-    # 生成可视化图
-    cell_diagram = draw_prokaryotic_cell_detailed(top_label)
-    
+    cell_diagram = draw_uniprot_style_cell(top_label)
     return confidences, cell_diagram
 
 # ==========================
-# 5. UI 界面构建 (Modern Scientific Theme)
+# 5. UI 界面 (Paper-White Style)
 # ==========================
-flagship_css = """
-@import url('https://fonts.googleapis.com/css2?family=IBM+Plex+Sans:wght@400;500;700&display=swap');
+# 极简学术风 CSS：白底、黑字、细线
+paper_css = """
+@import url('https://fonts.googleapis.com/css2?family=Roboto:wght@300;400;700&display=swap');
+body { font-family: 'Roboto', sans-serif !important; background-color: #ffffff; color: #1a1a1a; }
 
-body { font-family: 'IBM Plex Sans', sans-serif !important; background-color: #f8fafc; }
-
-/* 顶部 Header 样式 */
+/* 标题区：纯白背景，底部细线 */
 .header-box {
-    background: linear-gradient(145deg, #1e293b, #334155);
-    color: white; 
-    padding: 2.5rem; 
-    border-radius: 16px; 
+    background: #ffffff;
+    padding: 2rem 0;
+    border-bottom: 1px solid #e5e7eb;
     margin-bottom: 2rem;
-    box-shadow: 0 10px 25px -5px rgba(30, 41, 59, 0.3);
-    border: 1px solid rgba(255,255,255,0.1);
+    text-align: left;
 }
-.header-title { font-size: 2.4rem; font-weight: 700; letter-spacing: -0.02em; margin-bottom: 0.5rem; }
-.header-subtitle { opacity: 0.85; margin-bottom: 15px; font-size: 1.1rem; font-weight: 300; }
+.header-title { font-size: 2.5rem; font-weight: 700; color: #000000; letter-spacing: -1px; line-height: 1.1; }
+.header-subtitle { font-size: 1.2rem; color: #52525b; font-weight: 300; margin-top: 10px; }
 
-/* 徽章样式 */
-.badge-container { display: flex; gap: 10px; flex-wrap: wrap; }
+/* 标签徽章：简约线条风 */
 .badge { 
-    background: rgba(255,255,255,0.1); 
-    padding: 6px 14px; 
-    border-radius: 99px; 
-    font-size: 0.85rem; 
+    display: inline-block; 
+    padding: 4px 10px; 
+    font-size: 0.8rem; 
     font-weight: 500; 
-    border: 1px solid rgba(255,255,255,0.2); 
-    backdrop-filter: blur(4px); 
-}
-
-/* 内容卡片样式 */
-.content-box { 
-    background: white; 
-    padding: 2rem; 
-    border-radius: 16px; 
+    color: #0f172a; 
+    background: #f1f5f9; 
+    border-radius: 4px; 
     border: 1px solid #e2e8f0; 
-    box-shadow: 0 4px 6px -1px rgba(0,0,0,0.03), 0 2px 4px -1px rgba(0,0,0,0.02); 
-    height: 100%;
+    margin-right: 8px;
 }
 
-/* 图例说明框 */
-.legend-box {
-    margin-top:15px; 
-    padding:12px; 
-    background:#f1f5f9; 
-    border-radius:8px; 
-    font-size:0.85rem; 
-    color:#475569; 
-    border-left: 4px solid #64748b;
+/* 内容卡片：无阴影，仅边框，类似论文中的 Figure */
+.content-box {
+    background: #ffffff;
+    border: 1px solid #e5e7eb;
+    border-radius: 8px;
+    padding: 1.5rem;
 }
+
+/* 按钮：学术蓝 */
+button.primary { background: #2563eb !important; color: white !important; font-weight: 500; }
 """
 
-theme = gr.themes.Soft(
-    primary_hue="slate", 
-    font=[gr.themes.GoogleFont("IBM Plex Sans"), "ui-sans-serif", "system-ui"]
+theme = gr.themes.Base(
+    primary_hue="blue",
+    font=[gr.themes.GoogleFont("Roboto"), "ui-sans-serif", "system-ui"]
+).set(
+    body_background_fill="#ffffff",
+    block_background_fill="#ffffff",
+    block_border_width="1px",
+    block_label_background_fill="#ffffff"
 )
 
-with gr.Blocks(theme=theme, css=flagship_css, title="LocPred-Prok High-Fidelity") as app:
+with gr.Blocks(theme=theme, css=paper_css, title="LocPred-Prok") as app:
     
-    # --- Header Area ---
+    # --- Header (Paper Style) ---
     with gr.Column(elem_classes="header-box"):
         gr.HTML("""
             <div class="header-title">LocPred-Prok</div>
-            <div class="header-subtitle">High-Fidelity Prokaryotic Subcellular Localization Prediction</div>
-            <div class="badge-container">
-                <span class="badge">🧬 ESM-2 Enhanced</span>
-                <span class="badge">🔬 Scientific Visualization</span>
-                <span class="badge">🎯 Gram-negative Model</span>
+            <div class="header-subtitle">
+                Accurate prokaryotic subcellular localization using dual-branch protein language models
+            </div>
+            <div style="margin-top: 15px;">
+                <span class="badge">Article</span>
+                <span class="badge">ESM-2 Enhanced</span>
+                <span class="badge">Gram-negative</span>
             </div>
         """)
 
-    # --- Main Tab Area ---
     with gr.Tabs():
-        
-        # === Tab 1: Analysis ===
-        with gr.TabItem("🚀 Analysis", id="predict"):
+        with gr.TabItem("Prediction Interface"):
             with gr.Row():
-                
-                # Left Column: Input
-                with gr.Column(scale=5, elem_classes="content-box"):
-                    gr.Markdown("### 📥 Protein Sequence Input")
-                    sequence_input = gr.Textbox(
-                        lines=12, 
-                        placeholder=">Paste FASTA sequence here (e.g., >Protein ...)", 
-                        show_label=False
-                    )
+                # Input
+                with gr.Column(scale=4, elem_classes="content-box"):
+                    gr.Markdown("### 1. Sequence Input")
+                    gr.Markdown("<small style='color:gray'>Enter protein sequence in FASTA format.</small>")
+                    sequence_input = gr.Textbox(lines=10, show_label=False, placeholder=">Sequence...")
                     
                     with gr.Row():
-                        clear_btn = gr.ClearButton(sequence_input, value="Clear Input")
-                        submit_btn = gr.Button("Run Prediction & Visualization", variant="primary", scale=2)
+                        clear_btn = gr.ClearButton(sequence_input, value="Clear")
+                        submit_btn = gr.Button("Analyze", variant="primary")
                     
-                    gr.Markdown("#### 💡 Load Example Data")
+                    gr.Markdown("#### Example Data")
                     gr.Examples(
                         examples=[
-                            [">Outer Membrane Protein (OmpA)\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAFGGYQVNPYVGFEMGYDWLGRMPYKGSVENGAYKAQGVQLTAKLGYPITDDLDIYTRLGGMVWRADTKSNVYGKNHDTGVSPVFAGGVEYAITPEIATRLEYQWTNNIGDAHTIGTRPDNGMLSLGVSYRFGQGEAAPVVAPAPAPAPEVQTKHFTLKSDVLFNFNKATLKPEGQAALDQLYSQLSNLDPKDGSVVVLGYTDRIGSDAYNQGLSERRAQSVVDYLISKGIPADKISARGMGESNPVTGNTCDNVKQRAALIDCLAPDRRVEIEVKGIKDVVTQPQA"],
-                            [">Cell Wall Binding Protein\nMKFKLTAGCLAVAGVLLASSFGADAEIVVNAIYDQVARTEDGVYTQGQLTGRRIELLNKLGIEPEDSLASTVIHEFVARVGDDHGIETIIDEFYRQHPSASL"],
-                            [">Cytoplasmic Ribosomal Protein\nARYLGPKLKLSRREGTDLFLKSGVRAIDTKCKIEQAPGQHGARKPRLSDYGVQLREKQKVRRIYGVLERQFRNYYKEAARLKGNTGENLLALLEGRLDNVVYRMGFG"]
+                            [">Outer Membrane Protein A (OmpA)\nAPKNTWYTGAKLGWSQYHDTGFINNNGPTHENQLGAGAFGGYQVNPYVGFEMGYDWLGRMPYKGSVENGAYKAQGVQLTAKLGYPITDDLDIYTRLGGMVWRADTKSNVYGKNHDTGVSPVFAGGVEYAITPEIATRLEYQWTNNIGDAHTIGTRPDNGMLSLGVSYRFGQGEAAPVVAPAPAPAPEVQTKHFTLKSDVLFNFNKATLKPEGQAALDQLYSQLSNLDPKDGSVVVLGYTDRIGSDAYNQGLSERRAQSVVDYLISKGIPADKISARGMGESNPVTGNTCDNVKQRAALIDCLAPDRRVEIEVKGIKDVVTQPQA"],
+                            [">Cytoplasmic Protein (Ribo)\nARYLGPKLKLSRREGTDLFLKSGVRAIDTKCKIEQAPGQHGARKPRLSDYGVQLREKQKVRRIYGVLERQFRNYYKEAARLKGNTGENLLALLEGRLDNVVYRMGFG"]
                         ],
-                        inputs=sequence_input,
-                        label=None
+                        inputs=sequence_input, label=None
                     )
 
-                # Right Column: Output
-                with gr.Column(scale=7, elem_classes="content-box"):
-                    gr.Markdown("### 🔬 High-Fidelity Localization Map")
+                # Output
+                with gr.Column(scale=6, elem_classes="content-box"):
+                    gr.Markdown("### 2. Localization Analysis")
                     
-                    # 动态生成的细胞图
-                    output_image = gr.Image(
-                        label="Visualization", 
-                        show_label=False, 
-                        show_download_button=True, 
-                        interactive=False, 
-                        type="pil", 
-                        height=450
-                    )
+                    # 矢量风格图
+                    output_image = gr.Image(label="Schematic Visualization", show_label=False, show_download_button=True, interactive=False, type="pil", height=350)
                     
-                    gr.Markdown("### 📊 Prediction Confidence")
-                    # 概率分布条 (显示所有类别)
+                    # 概率分布
+                    gr.Markdown("#### Probability Scores")
                     output_label = gr.Label(num_top_classes=NUM_CLASSES, show_label=False)
 
-                    gr.HTML("""
-                        <div class="legend-box">
-                            <b>Legend Note:</b> This diagram illustrates a generic <b>Gram-negative bacterium</b> cross-section. 
-                            The peptidoglycan cell wall is depicted by the hatched mesh layer between the inner and outer membranes. 
-                            The nucleoid DNA is shown as tangled strands within the dense cytoplasm.
-                        </div>
-                    """)
-
-        # === Tab 2: About ===
-        with gr.TabItem("📖 Methodology"):
-             gr.Markdown("""
-             ### About LocPred-Prok
-             LocPred-Prok is a deep learning tool designed for precise prokaryotic protein subcellular localization.
-             
-             **Architecture:**
-             * **Backbone:** ESM-2 (150M Parameters)
-             * **Classifier:** Dual-Branch Architecture (Global Semantics + Local Motifs)
-             
-             **Visualization:**
-             The generated diagram dynamically renders a Gram-negative bacterial cell model using `matplotlib`, applying shading and texture effects to highlight the predicted localization site.
-             """)
-
-    # --- Event Handling ---
-    submit_btn.click(
-        fn=predict, 
-        inputs=sequence_input, 
-        outputs=[output_label, output_image]
-    )
-    clear_btn.click(
-        lambda: [None, None], 
-        outputs=[output_label, output_image]
-    )
-
-# 启动应用
-if __name__ == "__main__":
-    app.launch()
+        with gr.TabItem("About & Citation"):
+             gr.Markdown("### Method Description\nLocPred-Prok utilizes...")
+
+    submit_btn.click(fn=predict, inputs=sequence_input, outputs=[output_label, output_image])
+    clear_btn.click(lambda: [None, None], outputs=[output_label, output_image])
+
+app.launch()