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['text'])
+        self.structure=pdb_structure(dict_data['structure'])
     def __getitem__(self, index):
         return self.data['text'][index], self.structure[index]
     def __len__(self):
@@ -174,51 +174,57 @@ class MyModel(nn.Module):
         output_feature = self.dropout(self.relu(self.bn2(self.fc2(output_feature))))
         output_feature = self.dropout(self.relu(self.bn3(self.fc3(output_feature))))
         output_feature = self.dropout(self.output_layer(output_feature))
-        print(output_feature)
         return torch.softmax(output_feature,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)
+def pdb_structure(Structure_index):
+    created_folders = []
+    SurfacePoitCloud_all = []
+    for index in Structure_index:
+        structure_folder = join(temp_path, str(index))
+        os.makedirs(structure_folder, exist_ok=True)
+        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}")
+        coords, atname, pdbname, pdb_num = utils.parsePDB(structure_folder)
+        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)
+    for folder in created_folders:
+        shutil.rmtree(folder)
     return SurfacePoitCloud_all_tensor
 
 def ACE(file):
+    if not os.path.exists(pdb_path):
+        os.makedirs(pdb_path)
+    else:
+        shutil.rmtree(pdb_path)
+        os.makedirs(pdb_path)
     # df = pd.read_csv(seq_path)
     # test_sequences = df["Seq"].tolist()
     # test_Structure_index = df["Structure_index"].tolist()
         
     test_sequences = [file]
-    test_Structure_index = ["structure_1"]
+    test_Structure_index = [f"structure_{i}" for i in range(len(test_sequences))]
 
     
     test_dict = {"text":test_sequences, 'structure':test_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)
+    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)
     test_data=MyDataset(test_dict)
     test_dataloader=DataLoader(test_data,batch_size=batch_size,collate_fn=collate_fn,shuffle=False)
 
@@ -235,6 +241,7 @@ def ACE(file):
         print("=================================Start prediction========================")
         for index, (batch, structure_fea, fingerprint) in enumerate(test_dataloader):
             batchs = {k: v for k, v in batch.items()}
+            print(structure_fea)
             outputs = model(structure_fea, batchs, fingerprint)
             probability = outputs[0].tolist()
             print(outputs)
@@ -257,8 +264,12 @@ def ACE(file):
     summary['Probability'] = probability_all
     summary_df = pd.DataFrame(summary)
     summary_df.to_csv('output.csv', index=False)
-    out_text = output
-    out_prob = probability        
+    if len(test_sequences) > 1:
+        out_text = "Please download csv"
+        out_prob = "Please download csv"
+    else:
+        out_text = output
+        out_prob = probability        
     return 'output.csv', out_text, out_prob
 
 iface = gr.Interface(fn=ACE,