ESM2Bind / plot_pdb.py
wangjin2000's picture
Update plot_pdb.py
73969e8 verified
import py3Dmol
import requests
from pathlib import Path
def display_pdb_by_pdb(pdb):
# function to display pdb in py3dmol
# ref: https://huggingface.co/spaces/AIGE/A_B
view = py3Dmol.view(width=500, height=500)
view.setBackgroundColor('black');
view.addModel(pdb, "pdb")
view.setStyle({'cartoon': {'color': 'spectrum'}})
view.zoomTo()
output = view._make_html().replace("'", '"')
x = f"""<!DOCTYPE html><html></center> {output} </center></html>""" # do not use ' in this input
return f"""<iframe height="500px" width="100%" name="result" allow="midi; geolocation; microphone; camera;
display-capture; encrypted-media;" sandbox="allow-modals allow-forms
allow-scripts allow-same-origin allow-popups
allow-top-navigation-by-user-activation allow-downloads" allowfullscreen=""
allowpaymentrequest="" frameborder="0" srcdoc='{x}'></iframe>"""
def get_pdb(sequence):
retries = 0
pdb_str = None
url = "https://api.esmatlas.com/foldSequence/v1/pdb/"
while retries < 3 and pdb_str is None:
response = requests.post(url, data=sequence, verify=False)
pdb_str = response.text
if pdb_str == "INTERNAL SERVER ERROR":
retries += 1
time.sleep(0.1)
pdb = str = None
return pdb_str
def plot_struc(sequence):
headers = {
'Content-Type': 'application/x-www-form-urlencoded',
}
response = requests.post('https://api.esmatlas.com/foldSequence/v1/pdb/', headers=headers, data=sequence, verify=False) #verify=false jw 0425 work around for SSL certificate
pdb_string = get_pdb(sequence)
name = sequence[:3] + sequence[-3:]
outpath = (
Path.cwd() / f"PDB-{name}.pdb")
with open(outpath.name, "w") as f:
f.write(pdb_string)
outpath_str = str(outpath)
html_view = display_pdb_by_pdb(pdb_string)
return outpath_str, html_view