Update

.ipynb_checkpoints/app-checkpoint.py

@@ -42,6 +42,20 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
+reps = [
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "cartoon",
+        "color": "spectrum",
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True
+    }
+]
+
 def is_valid_sequence_length(length: int) -> bool:
     """Check if sequence length is within valid range."""
     return 100 <= length <= 1500
@@ -215,7 +229,6 @@ def fetch_pdb(pdb_id):
         print(f"Error fetching PDB: {e}")
         return None
 
-
 def score_to_color(score):
     norm = Normalize(vmin=0, vmax=1)  # Normalize scores between 0 and 1
     color_map = cm.coolwarm  # Directly use the colormap (e.g., 'cividis', 'coolwarm', etc.)
@@ -243,40 +256,9 @@ def process_pdb(pdb_id):
     # Prepare result string
     result_str = "\n".join([f"{aa}: {score:.2f}" for aa, score in zip(sequence, normalized_scores)])
 
-    # Prepare representations for Molecule3D
-    reps = [
-        {
-            "model": 0,
-            "chain": "",
-            "resname": "",
-            "style": "cartoon",
-            "color": "spectrum",
-            "residue_range": "",
-            "around": 0,
-            "byres": False,
-            "visible": True
-        }
-    ]
-
-    # Add color-coded residues based on binding site scores
-    #for i, score in enumerate(normalized_scores):
-    #    if score > 0.7:  # You can adjust this threshold
-    #        reps.append({
-    #            "model": 0,
-    #            "chain": chain.get_id(),
-    #            "style": "stick",
-    #            "color": score_to_color(score),
-    #            "residue_range": f"{i+1}-{i+1}",
-    #            "byres": True,
-    #            "visible": True
-    #        })
-
-    # Create Molecule3D with the PDB file and representations
-    molecule_viewer = Molecule3D(
-        reps=reps
-    )
-
-    return result_str, molecule_viewer
+    pdb_path = fetch_pdb(pdb_id)
+    
+    return result_str, pdb_path
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -292,10 +274,16 @@ with gr.Blocks() as demo:
             predict_btn = gr.Button("Predict Binding Sites")
 
         with gr.Column():
+            # Binding site predictions output
             predictions_output = gr.Textbox(
                 label="Binding Site Predictions"
             )
-            molecule_output = Molecule3D(label="Protein Structure")
+            
+            # 3D Molecule visualization
+            molecule_output = Molecule3D(
+                label="Protein Structure", 
+                reps=reps
+            )
 
     # Prediction logic
     predict_btn.click(

app.py

@@ -42,6 +42,20 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model.to(device)
 model.eval()
 
+reps = [
+    {
+        "model": 0,
+        "chain": "",
+        "resname": "",
+        "style": "cartoon",
+        "color": "spectrum",
+        "residue_range": "",
+        "around": 0,
+        "byres": False,
+        "visible": True
+    }
+]
+
 def is_valid_sequence_length(length: int) -> bool:
     """Check if sequence length is within valid range."""
     return 100 <= length <= 1500
@@ -215,7 +229,6 @@ def fetch_pdb(pdb_id):
         print(f"Error fetching PDB: {e}")
         return None
 
-
 def score_to_color(score):
     norm = Normalize(vmin=0, vmax=1)  # Normalize scores between 0 and 1
     color_map = cm.coolwarm  # Directly use the colormap (e.g., 'cividis', 'coolwarm', etc.)
@@ -243,40 +256,9 @@ def process_pdb(pdb_id):
     # Prepare result string
     result_str = "\n".join([f"{aa}: {score:.2f}" for aa, score in zip(sequence, normalized_scores)])
 
-    # Prepare representations for Molecule3D
-    reps = [
-        {
-            "model": 0,
-            "chain": "",
-            "resname": "",
-            "style": "cartoon",
-            "color": "spectrum",
-            "residue_range": "",
-            "around": 0,
-            "byres": False,
-            "visible": True
-        }
-    ]
-
-    # Add color-coded residues based on binding site scores
-    #for i, score in enumerate(normalized_scores):
-    #    if score > 0.7:  # You can adjust this threshold
-    #        reps.append({
-    #            "model": 0,
-    #            "chain": chain.get_id(),
-    #            "style": "stick",
-    #            "color": score_to_color(score),
-    #            "residue_range": f"{i+1}-{i+1}",
-    #            "byres": True,
-    #            "visible": True
-    #        })
-
-    # Create Molecule3D with the PDB file and representations
-    molecule_viewer = Molecule3D(
-        reps=reps
-    )
-
-    return result_str, molecule_viewer
+    pdb_path = fetch_pdb(pdb_id)
+    
+    return result_str, pdb_path
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -292,10 +274,16 @@ with gr.Blocks() as demo:
             predict_btn = gr.Button("Predict Binding Sites")
 
         with gr.Column():
+            # Binding site predictions output
             predictions_output = gr.Textbox(
                 label="Binding Site Predictions"
             )
-            molecule_output = Molecule3D(label="Protein Structure")
+            
+            # 3D Molecule visualization
+            molecule_output = Molecule3D(
+                label="Protein Structure", 
+                reps=reps
+            )
 
     # Prediction logic
     predict_btn.click(