add suport to cif

app.py

@@ -13,37 +13,10 @@ default_file = "static/test.pdb"
 TEMP_DIR = "static/tmp/"
 os.makedirs(TEMP_DIR, exist_ok=True)
 
-# # 保存全局信息以便交互
-# structure_cache = gr.State()
-# contact_cache = None
-
-def read_file(file_path):
-    if file_path is None:
-        return "<b style='color:red'>未提供结构文件</b>"
-
-    try:
-        with open(file_path, "r") as f:
-            structure_str = f.read()
-    except Exception as e:
-        return f"<b style='color:red'>读取文件失败: {e}</b>"
-
-    # file_format = file_path.split(".")[-1]
-    summary = analyze_structure_combined(file_path)
-    entity_color_dict = build_entity_color_dict(list(summary.keys()))
-
-    # 缓存用于后续交互
-    structure_dict = {
-        "structure_str": structure_str,
-        "summary": summary,
-        "entity_color_dict": entity_color_dict
-    }
-
-    return structure_str, summary, entity_color_dict, structure_dict
-
 
 def render_structure(structure_str, summary, entity_color_dict, add_label=True):
     view = py3Dmol.view(width=233, height=233)
-    view.addModel(structure_str, 'pdb')
+    view.addModel(structure_str) # 不指定类型似乎可以自动识别
     set_default_styles(view, summary, entity_color_dict, add_label=add_label)
     view.zoomTo()
     return view
@@ -174,7 +147,7 @@ with gr.Blocks() as demo:
                 )
                 pdb_btn = gr.Button("获取结构", scale=1)
 
-                file_input = gr.File(label="或直接上传PDB文件", file_types=[".pdb"])
+                file_input = gr.File(label="或直接上传PDB文件", file_types=[".pdb", ".cif", ".ent"])
             with gr.Column(scale=2):
                 with gr.Row():
                         entity_selector = gr.Dropdown(choices=[], interactive=True, multiselect=True, label="选择实体", scale=2)

utils.py

@@ -4,13 +4,13 @@ import hashlib
 import numpy as np
 
 from tempfile import NamedTemporaryFile
-
-from Bio.PDB import MMCIFParser, PDBParser, PDBList
+from io import StringIO
+from Bio.PDB import MMCIFParser, PDBParser, PDBIO
 from collections import defaultdict
 import colorsys
 import hashlib
 from Bio.PDB.NeighborSearch import NeighborSearch
-
+from Bio.PDB.DSSP import DSSP
 
 NUCLEIC_ACIDS = {
     "A", "G", "C", "U", "T",
@@ -112,7 +112,17 @@ def classify_residue(residue):
 
 def analyze_structure_combined(file_path):
     if file_path.endswith(".cif"):
+        with open(file_path, 'r') as f:
+            content = f.read()
+        # 如果缺少 data_ 开头，就加上一个默认块名
+        if not content.lstrip().startswith("data_"):
+            content = "data_auto\n" + content
+        # 3. 写入临时 mmCIF 文件
+        with NamedTemporaryFile(suffix=".cif", delete=False, mode='w') as tmp:
+            tmp.write(content)
+            file_path = tmp.name
         parser = MMCIFParser(QUIET=True)
+
     elif file_path.endswith(".pdb") or file_path.endswith(".ent"):
         parser = PDBParser(QUIET=True)
     else:
@@ -149,6 +159,30 @@ def analyze_structure_combined(file_path):
     ))
     return sorted_summary
 
+def read_file(file_path):
+    if file_path is None:
+        return "<b style='color:red'>未提供结构文件</b>"
+
+    try:
+        with open(file_path, "r") as f:
+            structure_str = f.read()
+    except Exception as e:
+        return f"<b style='color:red'>读取文件失败: {e}</b>"
+
+    # file_format = file_path.split(".")[-1]
+    summary = analyze_structure_combined(file_path)
+    entity_color_dict = build_entity_color_dict(list(summary.keys()))
+
+    # 缓存用于后续交互
+    structure_dict = {
+        "structure_str": structure_str,
+        "summary": summary,
+        "entity_color_dict": entity_color_dict
+    }
+
+    return structure_str, summary, entity_color_dict, structure_dict
+
+
 def extract_contact_residues(summary, selected_keys, cutoff=3.5):
     entity_atoms = {key: [] for key in selected_keys}
     atom_to_residue_info = {}
@@ -247,7 +281,7 @@ def extract_polar_contacts(summary, contact_summary, cutoff=3.5, angle_cutoff=12
 
 def set_default_styles(viewer, summary, entity_color_dict,
                        add_label=True):
-    # viewer.setStyle({'hetflag': True}, {"stick": {}})
+    viewer.setStyle({'hetflag': True}, {"stick": {}})
     for entity, color in entity_color_dict.items():
         label_style = {
                 'fontOpacity':1,
@@ -260,18 +294,19 @@ def set_default_styles(viewer, summary, entity_color_dict,
         if "(protein)" in entity.lower():
             # 提取链ID
             chain_id = entity.split()[1]
-            if chain_id == 'A':
-                print(summary[entity])
-            viewer.setStyle({'chain': chain_id}, {'cartoon': {'color': color}})
+            viewer.setStyle({'chain': chain_id},
+                            {'cartoon': {'arrows': True,
+                                         'color': color,
+                                         'opacity': 0.9}})
             if add_label:
                 viewer.addLabel(entity, label_style, {'chain':chain_id})
         elif '(dna)' in entity.lower() or '(rna)' in entity.lower():
             # 提取链ID
             chain_id = entity.split()[1]
-            if chain_id == 'T':
-                print(summary[entity])
             viewer.setStyle({'chain': chain_id},
-                            {'cartoon': {'color': 'blue', 'opacity': 0.8}})  # 碱基和磷酸用stick)
+                            {'cartoon': {'color': color,
+                                         'nucleicAcid': True,
+                                         'opacity': 0.8}})  # 碱基和磷酸用stick)
             if add_label:
                 viewer.addLabel(entity, label_style, {'chain':chain_id})