Update app.py

app.py

@@ -11,9 +11,6 @@ from Bio.Seq import Seq
 from Bio.SeqRecord import SeqRecord
 import time
 import urllib.parse
-import base64
-from IPython.display import Image
-import io
 
 urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
@@ -36,30 +33,7 @@ def render_mol(pdb):
     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()
-    
-    # Use the JavaScript method to get the PNG data as a base64 encoded string
-    png_base64 = pdbview.js().png()
-    
-    print(f"Type of png_base64: {type(png_base64)}")
-    print(f"Length of png_base64: {len(png_base64) if png_base64 else 'N/A'}")
-    
-    if png_base64 is None:
-        raise ValueError("Failed to capture image data")
-    
-    # If png_base64 is already a string, it's likely already base64 encoded
-    if isinstance(png_base64, str):
-        return png_base64
-    # If it's bytes, we need to encode it
-    elif isinstance(png_base64, bytes):
-        return base64.b64encode(png_base64).decode('utf-8')
-    else:
-        raise TypeError(f"Unexpected type for png_base64: {type(png_base64)}")
-        
 def perform_blast_analysis(sequence):
     st.subheader('Protein Analysis')
     with st.spinner("Analyzing generated protein... This may take a few minutes."):
@@ -175,14 +149,14 @@ if st.sidebar.button('Generate and Predict'):
 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'])
+    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)
+    col1, col2 = st.columns(2)
     with col1:
         if st.button('Analyze Protein'):
             perform_blast_analysis(st.session_state.sequence)
@@ -194,20 +168,11 @@ if st.session_state.structure_info:
             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("1. Take a screenshot of the protein structure 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.")
+    st.write("3. Attach your screenshot to the 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})")
@@ -220,7 +185,6 @@ if st.session_state.structure_info:
     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 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)!
     """)