File size: 1,726 Bytes
feb41b4 | 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 | import os
import sys
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
import argparse
from tqdm import tqdm
from structure.get_sst_seq import SSTPredictor
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--pdb_dir", type=str, default=None, help="Directory containing PDB files",)
parser.add_argument("--pdb_file", type=str, default=None, help="PDB file",)
parser.add_argument("--vocab_size", type=int, default=[2048], nargs='+', help="Vocabulary size",)
parser.add_argument("--overwrite", action="store_true", help="Overwrite existing files",)
parser.add_argument("--output_dir", type=str, default=None, help="Output directory",)
args = parser.parse_args()
os.makedirs(args.output_dir, exist_ok=True)
for v in args.vocab_size:
os.makedirs(os.path.join(args.output_dir, str(v)), exist_ok=True)
if args.pdb_dir is not None:
pdb_files = sorted(os.listdir(args.pdb_dir))
pdb_files = [os.path.join(args.pdb_dir, pdb_file) for pdb_file in pdb_files]
elif args.pdb_file is not None:
pdb_files = [args.pdb_file]
else:
raise ValueError("Either pdb_dir or pdb_file must be provided")
for v in args.vocab_size:
processor = SSTPredictor(structure_vocab_size=v)
results = processor.predict_from_pdb(pdb_files)
for result in results:
name = result['name'].split('.')[0]
sst_seq = result[f'{v}_sst_seq']
sst_seq = [str(i) for i in sst_seq]
with open(os.path.join(args.output_dir, str(v), name+'.fasta'), "w") as f:
f.write(f'>{name}\n')
f.write(','.join(sst_seq))
|