Spaces:
Sleeping
Sleeping
File size: 1,576 Bytes
b38c7b5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from typing import List
import __main__
import rootutils
import torch
from torch_geometric.data import Dataset
# setup root dir and pythonpath
rootutils.setup_root(__file__, indicator=".project-root", pythonpath=True)
from src.data.components.prepare_data import CropPairedPDB
setattr(__main__, "CropPairedPDB", CropPairedPDB)
class PinderDataset(Dataset):
"""Pinder dataset.
Args:
Dataset: PyTorch Geometric Dataset.
"""
def __init__(self, file_paths: List[str]) -> None:
"""Initialize the PinderDataset.
Args:
file_paths: List of file paths.
"""
super().__init__()
self.file_paths = file_paths
@property
def processed_file_names(self) -> List[str]:
"""Return the processed file names.
Returns:
List[str]: List of processed
"""
return self.file_paths
def len(self) -> int:
"""Return the length of the dataset.
Returns:
int: Length of the dataset
"""
return len(self.processed_file_names)
def get(self, idx) -> CropPairedPDB:
"""Get the data at the given index.
Args:
idx: Index of the data.
Returns:
CropPairedPDB: CropPairedPDB object.
"""
data = torch.load(self.processed_file_names[idx], weights_only=False)
return data
if __name__ == "__main__":
file_paths = ["./data/processed/apo/test/1a19__A1_P11540--1a19__B1_P11540.pt"]
dataset = PinderDataset(file_paths=file_paths)
print(dataset[0])
|