| 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)) |
| |
