Update app.py

app.py

@@ -42,16 +42,24 @@ def get_protein_image(pdbview):
     pdbview.zoom(0.8)
     pdbview.render()
     
-    # Get the image data directly as a base64 encoded string
+    # Use the JavaScript method to get the PNG data as a base64 encoded string
     png_base64 = pdbview.js().png()
     
-    # The png_base64 should already be a base64 encoded string
-    # If it's not, we might need to encode it
+    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
-    else:
+    # 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."):