Update app.py

app.py

@@ -51,7 +51,7 @@ class MyDataset(Dataset):
     def __init__(self,dict_data) -> None:
         super(MyDataset,self).__init__()
         self.data=dict_data
-        self.structure=pdb_structure(dict_data['structure'])
+        self.structure=pdb_structure(dict_data['test_sequences'])
     def __getitem__(self, index):
         return self.data['text'][index], self.structure[index]
     def __len__(self):
@@ -178,34 +178,29 @@ class MyModel(nn.Module):
         return torch.softmax(output_feature,dim=1)
 
 
-def pdb_structure(Structure_index):
-    created_folders = []
-    SurfacePoitCloud_all = []
-    for index in Structure_index:
-        # structure_folder = join(temp_path, str(index)+".pdb")
-        # created_folders.append(structure_folder)
-        # pdb_file = join(pdb_path, f"{index}.pdb")
-        # if os.path.exists(pdb_file):
-        #     shutil.copy2(pdb_file, structure_folder)
-        # else:
-        #     print(f"PDB file not found for structure {index}")
-        print("=================================Structure prediction========================")
-        for i in tqdm(range(0, len(test_sequences))):
-            command = ["curl", "-X", "POST", "-k", "--data", f"{test_sequences[i]}", "https://api.esmatlas.com/foldSequence/v1/pdb/"]
-            result = subprocess.run(command, capture_output=True, text=True)
-            # with open(os.path.join(pdb_path, f'{test_Structure_index[i]}.pdb'), 'w') as file:
-            #     file.write(result.stdout)
-        coords, atname, pdbname, pdb_num = utils.parsePDB(result.stdout)
-        atoms_channel = utils.atomlistToChannels(atname)
-        radius = utils.atomlistToRadius(atname)
-        PointCloudSurfaceObject = VolumeMaker.PointCloudSurface(device=device)
-        coords = coords.to(device)
-        radius = radius.to(device)
-        atoms_channel = atoms_channel.to(device)
-        SurfacePoitCloud = PointCloudSurfaceObject(coords, radius)
-        feature = SurfacePoitCloud.view(pdb_num,-1,3).cpu()
-        SurfacePoitCloud_all.append(feature)
-    SurfacePoitCloud_all_tensor = torch.squeeze(torch.stack(SurfacePoitCloud_all),dim=1)
+def pdb_structure(test_sequences):
+    # structure_folder = join(temp_path, str(index)+".pdb")
+    # created_folders.append(structure_folder)
+    # pdb_file = join(pdb_path, f"{index}.pdb")
+    # if os.path.exists(pdb_file):
+    #     shutil.copy2(pdb_file, structure_folder)
+    # else:
+    #     print(f"PDB file not found for structure {index}")
+    print("=================================Structure prediction========================")
+    command = ["curl", "-X", "POST", "-k", "--data", f"{test_sequences}", "https://api.esmatlas.com/foldSequence/v1/pdb/"]
+    result = subprocess.run(command, capture_output=True, text=True)
+    # with open(os.path.join(pdb_path, f'{test_Structure_index[i]}.pdb'), 'w') as file:
+    #     file.write(result.stdout)
+    coords, atname, pdbname, pdb_num = utils.parsePDB(result.stdout)
+    atoms_channel = utils.atomlistToChannels(atname)
+    radius = utils.atomlistToRadius(atname)
+    PointCloudSurfaceObject = VolumeMaker.PointCloudSurface(device=device)
+    coords = coords.to(device)
+    radius = radius.to(device)
+    atoms_channel = atoms_channel.to(device)
+    SurfacePoitCloud = PointCloudSurfaceObject(coords, radius)
+    feature = SurfacePoitCloud.view(pdb_num,-1,3).cpu()
+    SurfacePoitCloud_all_tensor = torch.squeeze(torch.stack(feature),dim=1)
     return SurfacePoitCloud_all_tensor
 
 def ACE(file):