esm3-conformational-sampling

Sleeping

App Files Files Community

Tonic commited on Oct 19, 2024

Commit

fbb545c

unverified ·

1 Parent(s): 9b5ec9d

attempt to fix the visualizer

Browse files

Files changed (1) hide show

app.py +18 -24

app.py CHANGED Viewed

@@ -236,17 +236,14 @@ def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
     if aligned is None:
         return None
-    viewer = py3Dmol.view(width=800, height=600)
-    viewer.addModel(protein_to_pdb(protein), "pdb")
-    viewer.setStyle({"cartoon": {"color": "lightgrey"}})
-    viewer.addModel(protein_to_pdb(aligned), "pdb")
-    viewer.setStyle({"model": -1}, {"cartoon": {"color": "lightgreen"}})
-    viewer.zoomTo()
-    return {
-        "pdb": protein_to_pdb(protein) + protein_to_pdb(aligned),
-        "style": {"cartoon": {"color": "spectrum"}}
-    }
 def protein_to_pdb(protein: ESMProtein):
     pdb_str = ""
@@ -289,28 +286,25 @@ def run_prediction(pdb_file, num_runs, noise_level, num_frames, progress=gr.Prog
             return "Please upload a PDB file.", "No file uploaded", "", ""
         progress(0, desc="Starting prediction")
-        view_data, crmsd_text = prediction_visualization(pdb_file, num_runs, noise_level, num_frames, progress)
         steric_clash_text, norm_steric_clash_text = calculate_clashes_for_pdb(pdb_file)
-        if view_data is None:
             return "No successful predictions were made. Try adjusting the parameters or check the PDB file.", crmsd_text, steric_clash_text, norm_steric_clash_text
         progress(0.9, desc="Rendering visualization")
         html_content = f"""
-        <div style="height: 600px; width: 100%;">
-            <script src="https://3dmol.csb.pitt.edu/build/3Dmol-min.js"></script>
-            <div id="container-{id(view_data)}" style="height: 100%; width: 100%; position: relative;"></div>
-            <script>
-                var viewer = $3Dmol.createViewer(document.getElementById("container-{id(view_data)}"), {{defaultcolors: $3Dmol.rasmolElementColors}});
-                viewer.addModel({json.dumps(view_data['pdb'])}, "pdb");
-                viewer.setStyle({{}}, {json.dumps(view_data['style'])});
-                viewer.zoomTo();
-                viewer.render();
-            </script>
-        </div>
         """
         progress(1.0, desc="Completed")
         return html_content, crmsd_text, steric_clash_text, norm_steric_clash_text
     except Exception as e:
         error_message = str(e)
         stack_trace = traceback.format_exc()

     if aligned is None:
         return None
+    view = py3Dmol.view(width=800, height=600)
+    view.addModel(protein_to_pdb(protein), "pdb")
+    view.setStyle({"cartoon": {"color": "lightgrey"}})
+    view.addModel(protein_to_pdb(aligned), "pdb")
+    view.setStyle({"model": -1}, {"cartoon": {"color": "lightgreen"}})
+    view.zoomTo()
+    return view
 def protein_to_pdb(protein: ESMProtein):
     pdb_str = ""
             return "Please upload a PDB file.", "No file uploaded", "", ""
         progress(0, desc="Starting prediction")
+        view, crmsd_text = prediction_visualization(pdb_file, num_runs, noise_level, num_frames, progress)
         steric_clash_text, norm_steric_clash_text = calculate_clashes_for_pdb(pdb_file)
+        if view is None:
             return "No successful predictions were made. Try adjusting the parameters or check the PDB file.", crmsd_text, steric_clash_text, norm_steric_clash_text
         progress(0.9, desc="Rendering visualization")
+        # Convert the py3Dmol view to HTML
+        view_html = view._make_html().replace("'", '"')
         html_content = f"""
+        <iframe style="width: 100%; height: 600px;" 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='<!DOCTYPE html><html>{view_html}</html>'></iframe>
         """
         progress(1.0, desc="Completed")
         return html_content, crmsd_text, steric_clash_text, norm_steric_clash_text
     except Exception as e:
         error_message = str(e)
         stack_trace = traceback.format_exc()