Update

.ipynb_checkpoints/app-checkpoint.py

@@ -184,42 +184,38 @@ def process_pdb(pdb_id):
     pdb_path = fetch_pdb(pdb_id)
     
     if not pdb_path:
-        return "Failed to fetch PDB file", None, None
-    
+        return "Failed to fetch PDB file", None
+
     # Extract protein sequence and chain
     protein_sequence, chain = extract_protein_sequence(pdb_path)
     
     if not protein_sequence:
-        return "No suitable protein sequence found", None, None
+        return "No suitable protein sequence found", None
     
     # Predict binding sites
     sequence, normalized_scores = predict_protein_sequence(protein_sequence)
     
     # Prepare representations for coloring residues
-    reps = []
-    for i, (res, score) in enumerate(zip(sequence, normalized_scores), start=1):
-        # Map score to a color gradient from blue (low) to red (high)
-        color_intensity = int(score * 255)
-        color = f'rgb({color_intensity}, 0, {255-color_intensity})'
-        
-        rep = {
+    reps = [
+        {
             "model": 0,
             "chain": chain.id,
             "resname": res,
-            "resnum": i,
+            "resnum": i+1,
             "style": "cartoon",
-            "color": color,
-            "residue_range": f"{i}-{i}",
+            "color": f'rgb({int(score * 255)}, 0, {int(255 - score * 255)})',
+            "residue_range": f"{i+1}-{i+1}",
             "around": 0,
             "byres": True,
             "visible": True
         }
-        reps.append(rep)
+        for i, (res, score) in enumerate(zip(sequence, normalized_scores))
+    ]
     
     # Prepare result string
     result_str = "\n".join([f"{aa}: {score:.2f}" for aa, score in zip(sequence, normalized_scores)])
     
-    return result_str, reps, pdb_path
+    return result_str, {"root": [pdb_path]}
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -253,7 +249,7 @@ with gr.Blocks() as demo:
     predict_btn.click(
         process_pdb, 
         inputs=[pdb_input], 
-        outputs=[predictions_output, molecule_output, molecule_output]
+        outputs=[predictions_output, molecule_output]
     )
 
     # Add some example inputs
@@ -265,7 +261,7 @@ with gr.Blocks() as demo:
             ["3TJN"]
         ],
         inputs=[pdb_input],
-        outputs=[predictions_output, molecule_output, molecule_output]
+        outputs=[predictions_output, molecule_output]
     )
 
 demo.launch(share=True)

app.py

@@ -184,42 +184,38 @@ def process_pdb(pdb_id):
     pdb_path = fetch_pdb(pdb_id)
     
     if not pdb_path:
-        return "Failed to fetch PDB file", None, None
-    
+        return "Failed to fetch PDB file", None
+
     # Extract protein sequence and chain
     protein_sequence, chain = extract_protein_sequence(pdb_path)
     
     if not protein_sequence:
-        return "No suitable protein sequence found", None, None
+        return "No suitable protein sequence found", None
     
     # Predict binding sites
     sequence, normalized_scores = predict_protein_sequence(protein_sequence)
     
     # Prepare representations for coloring residues
-    reps = []
-    for i, (res, score) in enumerate(zip(sequence, normalized_scores), start=1):
-        # Map score to a color gradient from blue (low) to red (high)
-        color_intensity = int(score * 255)
-        color = f'rgb({color_intensity}, 0, {255-color_intensity})'
-        
-        rep = {
+    reps = [
+        {
             "model": 0,
             "chain": chain.id,
             "resname": res,
-            "resnum": i,
+            "resnum": i+1,
             "style": "cartoon",
-            "color": color,
-            "residue_range": f"{i}-{i}",
+            "color": f'rgb({int(score * 255)}, 0, {int(255 - score * 255)})',
+            "residue_range": f"{i+1}-{i+1}",
             "around": 0,
             "byres": True,
             "visible": True
         }
-        reps.append(rep)
+        for i, (res, score) in enumerate(zip(sequence, normalized_scores))
+    ]
     
     # Prepare result string
     result_str = "\n".join([f"{aa}: {score:.2f}" for aa, score in zip(sequence, normalized_scores)])
     
-    return result_str, reps, pdb_path
+    return result_str, {"root": [pdb_path]}
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -253,7 +249,7 @@ with gr.Blocks() as demo:
     predict_btn.click(
         process_pdb, 
         inputs=[pdb_input], 
-        outputs=[predictions_output, molecule_output, molecule_output]
+        outputs=[predictions_output, molecule_output]
     )
 
     # Add some example inputs
@@ -265,7 +261,7 @@ with gr.Blocks() as demo:
             ["3TJN"]
         ],
         inputs=[pdb_input],
-        outputs=[predictions_output, molecule_output, molecule_output]
+        outputs=[predictions_output, molecule_output]
     )
 
 demo.launch(share=True)