fix visualiser

app.py

@@ -14,6 +14,8 @@ from esm.utils import residue_constants as RC
 import requests
 from dotenv import load_dotenv
 import torch
+import json 
+import time
 
 load_dotenv()
 
@@ -187,15 +189,16 @@ def align_after_prediction(protein: ESMProtein, structure_prediction: ESMProtein
     
 def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
     if aligned is None:
-        return py3Dmol.view(width=800, height=600)
+        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()
     
-    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
+    return viewer.render()
 
 def protein_to_pdb(protein: ESMProtein):
     pdb_str = ""
@@ -206,42 +209,58 @@ def protein_to_pdb(protein: ESMProtein):
                 pdb_str += f"ATOM  {i*37+j+1:5d}  {atom:3s} {aa:3s} A{i+1:4d}    {x:8.3f}{y:8.3f}{z:8.3f}\n"
     return pdb_str
 
-def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_frames: int):
+def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_frames: int, progress=gr.Progress()):
     protein = get_protein(pdb_file)
     runs = []
     
-    for frame in range(num_frames):
+    total_iterations = num_frames * num_runs
+    progress(0, desc="Starting predictions")
+    
+    for frame in progress.tqdm(range(num_frames), desc="Processing frames"):
         noisy_protein = add_noise_to_coordinates(protein, noise_level)
         
         for i in range(num_runs):
+            progress((frame * num_runs + i + 1) / total_iterations, desc=f"Frame {frame+1}, Run {i+1}")
             structure_prediction = run_structure_prediction(noisy_protein)
             if structure_prediction is not None:
                 aligned, crmsd = align_after_prediction(protein, structure_prediction)
                 if aligned is not None:
                     runs.append((crmsd, aligned))
+            time.sleep(0.1)  # Small delay to allow for UI updates
     
     if not runs:
         return None, "No successful predictions"
     
     best_aligned = sorted(runs, key=lambda x: x[0])[0]
-    view = visualize_after_pred(protein, best_aligned[1])
-    return view, f"Best cRMSD: {best_aligned[0]:.4f}"
+    view_data = visualize_after_pred(protein, best_aligned[1])
+    return view_data, f"Best cRMSD: {best_aligned[0]:.4f}"
 
-def run_prediction(pdb_file, num_runs, noise_level, num_frames):
+def run_prediction(pdb_file, num_runs, noise_level, num_frames, progress=gr.Progress()):
     try:
         if pdb_file is None:
             return "Please upload a PDB file.", "No file uploaded"
         
-        view, crmsd_text = prediction_visualization(pdb_file, num_runs, noise_level, num_frames)
-        if view is None:
+        progress(0, desc="Starting prediction")
+        view_data, crmsd_text = prediction_visualization(pdb_file, num_runs, noise_level, num_frames, progress)
+        if view_data is None:
             return "No successful predictions were made. Try adjusting the parameters or check the PDB file.", crmsd_text
         
-        html = view._make_html()
-        return f"""
-        <div style="height: 600px;">
-            {html}
+        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>
-        """, crmsd_text
+        """
+        progress(1.0, desc="Completed")
+        return html_content, crmsd_text
     except Exception as e:
         error_message = str(e)
         stack_trace = traceback.format_exc()