add steric clash results and fix visualization

app.py

@@ -243,7 +243,10 @@ def visualize_after_pred(protein: ESMProtein, aligned: ESMProtein):
     viewer.setStyle({"model": -1}, {"cartoon": {"color": "lightgreen"}})
     viewer.zoomTo()
     
-    return viewer.render()
+    return {
+        "pdb": protein_to_pdb(protein) + protein_to_pdb(aligned),
+        "style": {"cartoon": {"color": "spectrum"}}
+    }
 
 def protein_to_pdb(protein: ESMProtein):
     pdb_str = ""
@@ -283,13 +286,13 @@ def prediction_visualization(pdb_file, num_runs: int, noise_level: float, num_fr
 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"
+            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_clas_text = calculate_clashes_for_pdb(pdb_file)
+        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
+            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"""
@@ -306,7 +309,8 @@ def run_prediction(pdb_file, num_runs, noise_level, num_frames, progress=gr.Prog
         </div>
         """
         progress(1.0, desc="Completed")
-        return html_content, crmsd_text, steric_clash_text, norm_steric_clas_text
+        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()
@@ -317,8 +321,8 @@ def run_prediction(pdb_file, num_runs, noise_level, num_frames, progress=gr.Prog
             <h4>Stack Trace:</h4>
             <pre>{stack_trace}</pre>
         </div>
-        """, "Error occurred"
-
+        """, "Error occurred", "", ""
+    
 def create_demo():
     with gr.Blocks() as demo:
         gr.Markdown("# Protein Structure Prediction and Visualization with Noise and MD Frames")
@@ -334,13 +338,16 @@ def create_demo():
             with gr.Column(scale=2):
                 visualization = gr.HTML(label="3D Visualization")
                 alignment_result = gr.Textbox(label="Alignment Result")
+                steric_clash_result = gr.Textbox(label="Steric Clash Result")
+                norm_steric_clash_result = gr.Textbox(label="Normalized Steric Clash Result")
         
         run_button.click(
             fn=run_prediction,
             inputs=[pdb_file, num_runs, noise_level, num_frames],
-            outputs=[visualization, alignment_result]
+            outputs=[visualization, alignment_result, steric_clash_result, norm_steric_clash_result]
         )
 
+
         gr.Markdown("""
         ## How to use
         1. Upload a PDB file using the file uploader.