Commit ab8f4a5b authored by Geoffrey Yu's avatar Geoffrey Yu
Browse files

add ground truth feature processing

parent ca340ce8
...@@ -504,12 +504,14 @@ class OpenFoldSingleMultimerDataset(torch.utils.data.Dataset): ...@@ -504,12 +504,14 @@ class OpenFoldSingleMultimerDataset(torch.utils.data.Dataset):
data = self._parse_mmcif( data = self._parse_mmcif(
path, mmcif_id, chain, alignment_dir, alignment_index, path, mmcif_id, chain, alignment_dir, alignment_index,
) )
ground_truth.append(data) ground_truth_feats = self.feature_pipeline.process_features(data, self.mode)
ground_truth.append(ground_truth_feats)
elif(ext == ".core"): elif(ext == ".core"):
data = self.data_pipeline.process_core( data = self.data_pipeline.process_core(
path, alignment_dir, alignment_index, path, alignment_dir, alignment_index,
) )
ground_truth.append(data) ground_truth_feats = self.feature_pipeline.process_features(data, self.mode)
ground_truth.append(ground_truth_feats)
elif(ext == ".pdb"): elif(ext == ".pdb"):
structure_index = None structure_index = None
data = self.data_pipeline.process_pdb( data = self.data_pipeline.process_pdb(
...@@ -520,7 +522,8 @@ class OpenFoldSingleMultimerDataset(torch.utils.data.Dataset): ...@@ -520,7 +522,8 @@ class OpenFoldSingleMultimerDataset(torch.utils.data.Dataset):
alignment_index=alignment_index, alignment_index=alignment_index,
_structure_index=structure_index, _structure_index=structure_index,
) )
ground_truth.append(data) ground_truth_feats = self.feature_pipeline.process_features(data, self.mode)
ground_truth.append(ground_truth_feats)
else: else:
raise ValueError("Extension branch missing") raise ValueError("Extension branch missing")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment