Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+from stmol import showmol
 import py3Dmol
 import requests
 import biotite.structure.io as bsio
@@ -12,7 +13,6 @@ import time
 import urllib.parse
 import base64
 
-
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
 st.set_page_config(layout='wide')
@@ -26,14 +26,21 @@ def generate_sequence_from_words(words, length):
     return ''.join(random.choice(amino_acids) for _ in range(length))
 
 def render_mol(pdb):
-    pdbview = py3Dmol.view()
-    pdbview.addModel(pdb,'pdb')
-    pdbview.setStyle({'cartoon':{'color':'spectrum'}})
+    pdbview = py3Dmol.view(width=800, height=500)
+    pdbview.addModel(pdb, 'pdb')
+    pdbview.setStyle({'cartoon': {'color': 'spectrum'}})
     pdbview.setBackgroundColor('white')
     pdbview.zoomTo()
     pdbview.zoom(2, 800)
     pdbview.spin(True)
-    showmol(pdbview, height = 500,width=800)
+    showmol(pdbview, height=500, width=800)
+    return pdbview
+
+def get_protein_image(pdbview):
+    pdbview.zoom(0.8)
+    pdbview.render()
+    png = pdbview.pngImage()
+    return base64.b64encode(png).decode('utf-8')
 
 def perform_blast_analysis(sequence):
     st.subheader('Protein Analysis')
@@ -113,6 +120,11 @@ def update(sequence, word1, word2, word3, sequence_length):
         st.error(f"An error occurred while calling the API: {str(e)}")
         st.write("Please try again later or contact support if the issue persists.")
 
+def share_on_twitter(word1, word2, word3, length, plddt):
+    tweet_text = f"I generated a unique protein structure from the words '{word1}', '{word2}', and '{word3}' with length {length}! plDDT Score: {plddt}% Try it yourself at [Your App URL] #GenPro2 #ProteinFolding"
+    tweet_url = f"https://twitter.com/intent/tweet?text={urllib.parse.quote(tweet_text)}"
+    return tweet_url
+
 # Initialize session state variables
 if 'sequence' not in st.session_state:
     st.session_state.sequence = None
@@ -142,45 +154,6 @@ if st.sidebar.button('Generate and Predict'):
         st.sidebar.warning("Please enter all three words to generate a sequence.")
 
 # Display structure information if available
-if st.session_state.structure_info:
-    info = st.session_state.structure_info
-    st.subheader(f'Predicted protein structure using seed: {info["word1"]}, {info["word2"]}, and {info["word3"]} + length {info["sequence_length"]}')
-    render_mol(info['pdb_string'])
-    
-    st.subheader('plDDT Score')
-    st.write('plDDT is a per-residue estimate of the confidence in prediction on a scale from 0-100%.')
-    st.info(f'Average plDDT: {int(info["b_value"] * 100)}%')
-    
-    st.download_button(
-        label="Download PDB",
-        data=info['pdb_string'],
-        file_name='predicted.pdb',
-        mime='text/plain',
-    )
-
-def render_mol(pdb):
-    pdbview = py3Dmol.view(width=800, height=500)
-    pdbview.addModel(pdb, 'pdb')
-    pdbview.setStyle({'cartoon': {'color': 'spectrum'}})
-    pdbview.setBackgroundColor('white')
-    pdbview.zoomTo()
-    pdbview.zoom(2, 800)
-    pdbview.spin(True)
-    showmol(pdbview, height=500, width=800)
-    return pdbview
-
-def get_protein_image(pdbview):
-    pdbview.zoom(0.8)
-    pdbview.render()
-    png = pdbview.pngImage()
-    return base64.b64encode(png).decode('utf-8')
-
-def share_on_twitter(word1, word2, word3, length, plddt):
-    tweet_text = f"I generated a unique protein structure from the words '{word1}', '{word2}', and '{word3}' with length {length}! plDDT Score: {plddt}% Try it yourself at [Your App URL] #GenPro2 #ProteinFolding"
-    tweet_url = f"https://twitter.com/intent/tweet?text={urllib.parse.quote(tweet_text)}"
-    return tweet_url
-
-# In the main app code, modify the structure display section:
 if st.session_state.structure_info:
     info = st.session_state.structure_info
     st.subheader(f'Predicted protein structure using seed: {info["word1"]}, {info["word2"]}, and {info["word3"]} + length {info["sequence_length"]}')
@@ -221,7 +194,6 @@ if st.session_state.structure_info:
     tweet_url = share_on_twitter(info["word1"], info["word2"], info["word3"], info["sequence_length"], plddt_score)
     st.markdown(f"[Share Results]({tweet_url})")
 
-    
     st.markdown("""
     ## What to do next:
     If you find interesting results from the sequence folding, you can explore further:
@@ -229,13 +201,8 @@ if st.session_state.structure_info:
     2. Visit the [Protein Data Bank (PDB)](https://www.rcsb.org/) for known protein structures.
     3. Compare your folded structure with known functional proteins by downloading your results.
     4. Read about similar proteins to gain insights into potential functions.
-    5. Click the "Analyze Protein" button below to get more information about your generated protein.
+    5. Click the "Analyze Protein" button to get more information about your generated protein.
 
     **Remember, this folding is based on randomly generated sequences. Interpret the results with caution.
     Enjoy exploring the world of protein sequences! Share your high-confidence protein images with us on X [*@WandsAI*](https://x.com/wandsai)!
-    """)
-
-# Show the Analyze Protein button if a sequence has been generated
-if st.session_state.show_analyze_button:
-    if st.button('Analyze Protein'):
-        perform_blast_analysis(st.session_state.sequence)
+    """)