Spaces:
Sleeping
Sleeping
File size: 2,890 Bytes
869dc68 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 |
import os
import contextlib
import torch
import warnings
from tqdm import tqdm, TqdmExperimentalWarning
from functools import partialmethod
import streamlit as st
from stmol import *
from chroma.utility.api import register_key
from dotenv import load_dotenv
import locale
from chroma import Chroma, Protein, conditioners
from chroma.models import graph_classifier, procap
from chroma.utility.chroma import letter_to_point_cloud, plane_split_protein
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
warnings.filterwarnings("ignore", category=TqdmExperimentalWarning)
tqdm.__init__ = partialmethod(tqdm.__init__, leave=False)
load_dotenv()
api_key = os.getenv("API_KEY")
register_key(api_key)
locale.getpreferredencoding = lambda: "UTF-8"
device = 'cuda' if torch.cuda.is_available() else 'cpu'
with contextlib.redirect_stdout(None):
chroma = Chroma(device=device)
def download(outputFile,newFileName,description):
with open(outputFile, "rb") as file:
btn = st.download_button(
label=description,
data=file,
file_name=newFileName,
)
def display(output,style,resn):
# imformation
protein=Protein.from_PDB(output,device=device)
st.subheader("Protein Information:")
st.write(f"Device: GPU")
st.write(f"Protein Length: {len(protein)} residues")
st.write(f"Structured Residue Count: {protein.length(structured=True)}")
# 显示 Protein 的序列
st.subheader("Protein Sequence:")
protein_sequence = protein.sequence(format="three-letter-list")
st.markdown(f"**Protein Sequence:** {protein_sequence}")
st.write(protein_sequence)
# 显示 Protein 的结构
with open(output, "r") as file:
pdb_content = file.read()
obj = makeobj(pdb_content,style=style,background='white')
# 使用 stmol 展示蛋白质结构
st.subheader("Protein Structure:")
traj_output = output.replace(".pdb", "_trajectory.pdb")
protein_newName = st.text_input("The specified file name. Default is {}.".format(output[output.rfind("/") + 1:])+"Please press [Enter] to confirm the change before download.", value=output[output.rfind("/") + 1:], key='protein_newName')
download(output,protein_newName,"Download sample")
traj_newName = st.text_input("The specified file name. Default is {}.".format(traj_output[traj_output.rfind("/") + 1:])+"Please press [Enter] to confirm the change before download.", value=traj_output[traj_output.rfind("/") + 1:], key='traj_newName')
download(traj_output,traj_newName,"Download trajectory")
if resn !='*':
obj = render_pdb_resn(obj ,resn_lst =resn)
showmol(obj, width=1800)
def render(protein, trajectories, output="./output/protein.pdb"):
protein.to_PDB(output)
traj_output = output.replace(".pdb", "_trajectory.pdb")
trajectories["trajectory"].to_PDB(traj_output)
|