fix infewrence_app.py

inference_app.py

@@ -108,6 +108,33 @@ from pinder.core.structure.atoms import atom_array_from_pdb_file
 from pathlib import Path
 from pinder.eval.dockq.biotite_dockq import BiotiteDockQ
 
+def extract_coordinates_from_pdb(filename, atom_name="CA"):
+    """
+    Extracts coordinates for specific atoms from a PDB file and returns them as a list of tuples.
+    Each tuple contains (x, y, z) coordinates of the specified atom type.
+    
+    Parameters:
+        filename (str): Path to the PDB file.
+        atom_name (str): The name of the atom to filter by (e.g., "CA" for alpha carbon).
+        
+    Returns:
+        list of tuple: List of coordinates as (x, y, z) tuples for the specified atom.
+    """
+    parser = PDB.PDBParser(QUIET=True)
+    structure = parser.get_structure("structure", filename)
+
+    coordinates = []
+
+    # Loop through each model, chain, residue, and atom to collect coordinates of specified atom
+    for model in structure:
+        for chain in model:
+            for residue in chain:
+                for atom in residue:
+                      # Filter for specific atom name
+                        xyz = atom.coord  # Coordinates are in a numpy array
+                        coordinates.append([xyz[0], xyz[1], xyz[2]])
+
+    return coordinates
 
 log = setup_logger(__name__)
 
@@ -285,8 +312,8 @@ def create_graph(pdb1, pdb2, k=5):
         HeteroData: A PyG HeteroData object containing ligand and receptor data.
     """
     # Extract coordinates from PDB files
-    coords1 = torch.tensor(atom_array_from_pdb_file(pdb1),dtype=torch.float)  
-    coords2 = torch.tensor(atom_array_from_pdb_file(pdb2),dtype=torch.float)  
+    coords1 = torch.tensor(extract_coordinates_from_pdb(pdb1),dtype=torch.float)  
+    coords2 = torch.tensor(extract_coordinates_from_pdb(pdb2),dtype=torch.float)  
     # coords3 = torch.tensor(extract_coordinates_from_pdb(pdb3),dtype=torch.float)
     # Create the HeteroData object
     data = HeteroData()
@@ -881,7 +908,7 @@ def predict (input_seq_1, input_msa_1, input_protein_1, input_seq_2,input_msa_2,
         mat, vect = model(data)
     mat = mat.to(device)
     vect = vect.to(device)
-    ligand1 = torch.tensor(atom_array_from_pdb_file(input_protein_1),dtype=torch.float).to(device)
+    ligand1 = torch.tensor(extract_coordinates_from_pdb(input_protein_1),dtype=torch.float).to(device)
     # receptor1 = torch.tensor(extract_coordinates_from_pdb(input_protein_2),dtype=torch.float).to(device)
     transformed_ligand = torch.matmul(ligand1, mat) + vect
     # transformed_receptor = torch.matmul(receptor1, mat) + vect