(self, idx)
| 232 | return data_dict, surf_dict |
| 233 | |
| 234 | def get_item(self, idx): |
| 235 | with open(self.pkl_files[idx], "rb") as fin: |
| 236 | data_dict = pickle.load(fin) |
| 237 | if self.surf_path: |
| 238 | surf_file = os.path.join(self.surf_path, os.path.basename(self.pkl_files[idx])) |
| 239 | with open(surf_file, "rb") as fin: |
| 240 | surf_dict = pickle.load(fin) |
| 241 | else: |
| 242 | surf_dict = None |
| 243 | if self.max_length: |
| 244 | data_dict, surf_dict = self.truncate(data_dict, surf_dict=surf_dict) |
| 245 | protein = load_protein(data_dict) |
| 246 | |
| 247 | item = {"graph": protein} |
| 248 | if surf_dict is not None: |
| 249 | surf_graph = load_surface(surf_dict) |
| 250 | item.update({"surf_graph": surf_graph}) |
| 251 | with protein.residue(): |
| 252 | protein.res2surf = torch.as_tensor(surf_dict["res2surf"]) # Need to transform local index to global index after batching |
| 253 | if self.transform: |
| 254 | item = self.transform(item) |
| 255 | return item |
| 256 | |
| 257 | def __len__(self): |
| 258 | return len(self.pkl_files) |
nothing calls this directly
no test coverage detected