Update

.ipynb_checkpoints/app-checkpoint.py

@@ -177,45 +177,28 @@ def extract_protein_sequence(pdb_path):
                 longest_sequence = one_letter_sequence
                 longest_chain = chain
     
-    return longest_sequence, longest_chain
+    return longest_sequence, longest_chain, pdb_path
 
 def process_pdb(pdb_id):
     # Fetch PDB file
     pdb_path = fetch_pdb(pdb_id)
     
     if not pdb_path:
-        return "Failed to fetch PDB file", None
+        return "Failed to fetch PDB file", pdb_path
 
     # Extract protein sequence and chain
-    protein_sequence, chain = extract_protein_sequence(pdb_path)
+    protein_sequence, chain, pdb_file = extract_protein_sequence(pdb_path)
     
     if not protein_sequence:
-        return "No suitable protein sequence found", None
-    
+        return "No suitable protein sequence found", pdb_file
+
     # Predict binding sites
     sequence, normalized_scores = predict_protein_sequence(protein_sequence)
     
-    # Prepare representations for coloring residues
-    reps = [
-        {
-            "model": 0,
-            "chain": chain.id,
-            "resname": res,
-            "resnum": i+1,
-            "style": "cartoon",
-            "color": f'rgb({int(score * 255)}, 0, {int(255 - score * 255)})',
-            "residue_range": f"{i+1}-{i+1}",
-            "around": 0,
-            "byres": True,
-            "visible": True
-        }
-        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, {"root": [pdb_path]}
+    return result_str, pdb_file
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -241,8 +224,7 @@ with gr.Blocks() as demo:
             
             # 3D Molecule visualization
             molecule_output = Molecule3D(
-                label="Protein Structure", 
-                reps=[]  # Start with empty representations
+                label="Protein Structure"
             )
     
     # Prediction logic

app.py

@@ -177,45 +177,28 @@ def extract_protein_sequence(pdb_path):
                 longest_sequence = one_letter_sequence
                 longest_chain = chain
     
-    return longest_sequence, longest_chain
+    return longest_sequence, longest_chain, pdb_path
 
 def process_pdb(pdb_id):
     # Fetch PDB file
     pdb_path = fetch_pdb(pdb_id)
     
     if not pdb_path:
-        return "Failed to fetch PDB file", None
+        return "Failed to fetch PDB file", pdb_path
 
     # Extract protein sequence and chain
-    protein_sequence, chain = extract_protein_sequence(pdb_path)
+    protein_sequence, chain, pdb_file = extract_protein_sequence(pdb_path)
     
     if not protein_sequence:
-        return "No suitable protein sequence found", None
-    
+        return "No suitable protein sequence found", pdb_file
+
     # Predict binding sites
     sequence, normalized_scores = predict_protein_sequence(protein_sequence)
     
-    # Prepare representations for coloring residues
-    reps = [
-        {
-            "model": 0,
-            "chain": chain.id,
-            "resname": res,
-            "resnum": i+1,
-            "style": "cartoon",
-            "color": f'rgb({int(score * 255)}, 0, {int(255 - score * 255)})',
-            "residue_range": f"{i+1}-{i+1}",
-            "around": 0,
-            "byres": True,
-            "visible": True
-        }
-        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, {"root": [pdb_path]}
+    return result_str, pdb_file
 
 # Create Gradio interface
 with gr.Blocks() as demo:
@@ -241,8 +224,7 @@ with gr.Blocks() as demo:
             
             # 3D Molecule visualization
             molecule_output = Molecule3D(
-                label="Protein Structure", 
-                reps=[]  # Start with empty representations
+                label="Protein Structure"
             )
     
     # Prediction logic