Update

.ipynb_checkpoints/app-checkpoint.py

@@ -22,11 +22,16 @@ from scipy.special import expit
 
 import requests
 
+from gradio_molecule3d import Molecule3D
+
 # Biopython imports
 from Bio.PDB import PDBParser, Select, PDBIO
 from Bio.PDB.DSSP import DSSP
 import Bio.PDB.PDBList as PDBList
 
+from matplotlib import cm  # For color mapping
+from matplotlib.colors import Normalize
+
 # Configuration
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
 max_length = 1500
@@ -210,64 +215,77 @@ def fetch_pdb(pdb_id):
         print(f"Error fetching PDB: {e}")
         return None
 
+# Function to map scores to colors (blue for low scores, red for high scores)
+def score_to_color(score):
+    norm = Normalize(vmin=0, vmax=1)  # Assuming scores are normalized between 0 and 1
+    color_map = cm.get_cmap('coolwarm')  # Use a blue-to-red colormap
+    rgba = color_map(norm(score))  # Get RGBA values
+    hex_color = '#{:02x}{:02x}{:02x}'.format(int(rgba[0] * 255), int(rgba[1] * 255), int(rgba[2] * 255))
+    return hex_color
+
 def process_pdb(pdb_id):
     # Fetch PDB file
-    # Use PDBList to download the file if it doesn't exist locally
     pdbl = PDBList.PDBList()
     pdb_path = pdbl.retrieve_pdb_file(pdb_id, pdir='pdb_files', file_format='pdb')
-    
+
     if not pdb_path or not os.path.exists(pdb_path):
         return "Failed to fetch PDB file", None
 
     # Extract protein sequence and chain
     protein_sequence, chain, filtered_pdb_path = extract_protein_sequence(pdb_path)
-    
+
     if not protein_sequence:
         return "No suitable protein sequence found", None
 
     # Predict binding sites
     sequence, normalized_scores = predict_protein_sequence(protein_sequence)
-    
+
     # Prepare result string
     result_str = "\n".join([f"{aa}: {score:.2f}" for aa, score in zip(sequence, normalized_scores)])
-    
-    return result_str, filtered_pdb_path
+
+    # Prepare residue-based coloring for Molecule3D
+    reps = []
+    for i, score in enumerate(normalized_scores):
+        reps.append({
+            "model": 0,
+            "chain": chain.get_id(),
+            "residue_range": f"{i}-{i}",
+            "style": "stick",
+            "color": score_to_color(score),
+            "byres": True,
+            "visible": True
+        })
+
+    molecule_viewer = Molecule3D(reps=reps)
+
+    return result_str, molecule_viewer
 
 # Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Protein Binding Site Prediction")
-    
+
     with gr.Row():
         with gr.Column():
-            # PDB ID input with default suggestion
             pdb_input = gr.Textbox(
-                value="2IWI", 
+                value="2IWI",
                 label="PDB ID",
                 placeholder="Enter PDB ID here..."
             )
-            
-            # Predict button
             predict_btn = gr.Button("Predict Binding Sites")
-        
+
         with gr.Column():
-            # Binding site predictions output
             predictions_output = gr.Textbox(
                 label="Binding Site Predictions"
             )
-            
-            # 3D Molecule visualization
-            molecule_output = Molecule3D(
-                label="Protein Structure"
-            )
-    
+            molecule_output = Molecule3D(label="Protein Structure")
+
     # Prediction logic
     predict_btn.click(
-        process_pdb, 
-        inputs=[pdb_input], 
+        process_pdb,
+        inputs=[pdb_input],
         outputs=[predictions_output, molecule_output]
     )
 
-    # Add some example inputs
     gr.Markdown("## Examples")
     gr.Examples(
         examples=[

.ipynb_checkpoints/requirements-checkpoint.txt

@@ -10,4 +10,5 @@ sentencepiece
 huggingface_hub>=0.15.0
 requests
 gradio_molecule3d
-biopython>=1.81
+biopython>=1.81
+matplotlib

app.py

@@ -22,11 +22,16 @@ from scipy.special import expit
 
 import requests
 
+from gradio_molecule3d import Molecule3D
+
 # Biopython imports
 from Bio.PDB import PDBParser, Select, PDBIO
 from Bio.PDB.DSSP import DSSP
 import Bio.PDB.PDBList as PDBList
 
+from matplotlib import cm  # For color mapping
+from matplotlib.colors import Normalize
+
 # Configuration
 checkpoint = 'ThorbenF/prot_t5_xl_uniref50'
 max_length = 1500
@@ -210,64 +215,77 @@ def fetch_pdb(pdb_id):
         print(f"Error fetching PDB: {e}")
         return None
 
+# Function to map scores to colors (blue for low scores, red for high scores)
+def score_to_color(score):
+    norm = Normalize(vmin=0, vmax=1)  # Assuming scores are normalized between 0 and 1
+    color_map = cm.get_cmap('coolwarm')  # Use a blue-to-red colormap
+    rgba = color_map(norm(score))  # Get RGBA values
+    hex_color = '#{:02x}{:02x}{:02x}'.format(int(rgba[0] * 255), int(rgba[1] * 255), int(rgba[2] * 255))
+    return hex_color
+
 def process_pdb(pdb_id):
     # Fetch PDB file
-    # Use PDBList to download the file if it doesn't exist locally
     pdbl = PDBList.PDBList()
     pdb_path = pdbl.retrieve_pdb_file(pdb_id, pdir='pdb_files', file_format='pdb')
-    
+
     if not pdb_path or not os.path.exists(pdb_path):
         return "Failed to fetch PDB file", None
 
     # Extract protein sequence and chain
     protein_sequence, chain, filtered_pdb_path = extract_protein_sequence(pdb_path)
-    
+
     if not protein_sequence:
         return "No suitable protein sequence found", None
 
     # Predict binding sites
     sequence, normalized_scores = predict_protein_sequence(protein_sequence)
-    
+
     # Prepare result string
     result_str = "\n".join([f"{aa}: {score:.2f}" for aa, score in zip(sequence, normalized_scores)])
-    
-    return result_str, filtered_pdb_path
+
+    # Prepare residue-based coloring for Molecule3D
+    reps = []
+    for i, score in enumerate(normalized_scores):
+        reps.append({
+            "model": 0,
+            "chain": chain.get_id(),
+            "residue_range": f"{i}-{i}",
+            "style": "stick",
+            "color": score_to_color(score),
+            "byres": True,
+            "visible": True
+        })
+
+    molecule_viewer = Molecule3D(reps=reps)
+
+    return result_str, molecule_viewer
 
 # Create Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Protein Binding Site Prediction")
-    
+
     with gr.Row():
         with gr.Column():
-            # PDB ID input with default suggestion
             pdb_input = gr.Textbox(
-                value="2IWI", 
+                value="2IWI",
                 label="PDB ID",
                 placeholder="Enter PDB ID here..."
             )
-            
-            # Predict button
             predict_btn = gr.Button("Predict Binding Sites")
-        
+
         with gr.Column():
-            # Binding site predictions output
             predictions_output = gr.Textbox(
                 label="Binding Site Predictions"
             )
-            
-            # 3D Molecule visualization
-            molecule_output = Molecule3D(
-                label="Protein Structure"
-            )
-    
+            molecule_output = Molecule3D(label="Protein Structure")
+
     # Prediction logic
     predict_btn.click(
-        process_pdb, 
-        inputs=[pdb_input], 
+        process_pdb,
+        inputs=[pdb_input],
         outputs=[predictions_output, molecule_output]
     )
 
-    # Add some example inputs
     gr.Markdown("## Examples")
     gr.Examples(
         examples=[

requirements.txt

@@ -10,4 +10,5 @@ sentencepiece
 huggingface_hub>=0.15.0
 requests
 gradio_molecule3d
-biopython>=1.81
+biopython>=1.81
+matplotlib