Update app.py

app.py

@@ -1,5 +1,4 @@
 import streamlit as st
-from stmol import showmol
 import py3Dmol
 import requests
 import biotite.structure.io as bsio
@@ -10,6 +9,9 @@ from Bio.Blast import NCBIWWW, NCBIXML
 from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 import time
+import urllib.parse
+import base64
+
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
@@ -156,6 +158,70 @@ if st.session_state.structure_info:
         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"]}')
+    pdbview = 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%.')
+    plddt_score = int(info["b_value"] * 100)
+    st.info(f'Average plDDT: {plddt_score}%')
+    
+    col1, col2, col3 = st.columns(3)
+    with col1:
+        if st.button('Analyze Protein'):
+            perform_blast_analysis(st.session_state.sequence)
+    
+    with col2:
+        st.download_button(
+            label="Download PDB",
+            data=info['pdb_string'],
+            file_name='predicted.pdb',
+            mime='text/plain',
+        )
+    
+    with col3:
+        protein_image = get_protein_image(pdbview)
+        st.download_button(
+            label="Download Image",
+            data=base64.b64decode(protein_image),
+            file_name="protein_structure.png",
+            mime="image/png"
+        )
+
+    st.subheader("Share your unique protein on X")
+    st.write("1. Download the protein image using the 'Download Image' button above.")
+    st.write("2. Click the 'Share Results' button below to open a pre-filled tweet.")
+    st.write("3. Attach the downloaded image to your tweet before posting.")
+    
+    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: