| import os |
| import shutil |
| import argparse |
| from tqdm import tqdm |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description='data format converter') |
| parser.add_argument("--input_dir", type=str, default=None, required=True) |
| parser.add_argument("--output_dir", type=str, default=None, required=True) |
| parser.add_argument("--format", type=str, default="REM2SSN", choices=['SSN2REM', 'REM2SSN'], required=True) |
| args = parser.parse_args() |
| |
| os.makedirs(args.output_dir, exist_ok=True) |
| if args.format == "REM2SSN": |
| output_dir = os.path.join(args.output_dir, "DATASET") |
| os.makedirs(output_dir, exist_ok=True) |
| proteins = os.listdir(args.input_dir + "/aa_seq") |
| for p in tqdm(proteins): |
| protein_name = p.split(".")[0] |
| out_protein_dir = os.path.join(output_dir, protein_name) |
| os.makedirs(out_protein_dir, exist_ok=True) |
| shutil.copyfile(f"{args.input_dir}/aa_seq/{protein_name}.fasta", f"{out_protein_dir}/{protein_name}.fasta") |
| shutil.copyfile(f"{args.input_dir}/pdbs/{protein_name}.pdb", f"{out_protein_dir}/{protein_name}.pdb") |
| shutil.copyfile(f"{args.input_dir}/substitutions/{protein_name}.csv", f"{out_protein_dir}/{protein_name}.csv") |
| elif args.format == "SSN2REM": |
| proteins = os.listdir(args.input_dir + "/DATASET") |
| pdb_dir = os.path.join(args.output_dir, "pdbs") |
| os.makedirs(pdb_dir, exist_ok=True) |
| aa_seq_dir = os.path.join(args.output_dir, "aa_seq") |
| os.makedirs(aa_seq_dir, exist_ok=True) |
| substitutions_dir = os.path.join(args.output_dir, "substitutions") |
| os.makedirs(substitutions_dir, exist_ok=True) |
| for p in tqdm(proteins): |
| protein_name = p |
| out_protein_dir = os.path.join(args.output_dir, protein_name) |
| os.makedirs(out_protein_dir, exist_ok=True) |
| shutil.copyfile(f"{args.input_dir}/DATASET/{protein_name}/{protein_name}.fasta", f"{out_protein_dir}/{protein_name}.fasta") |
| shutil.copyfile(f"{args.input_dir}/DATASET/{protein_name}/{protein_name}.pdb", f"{out_protein_dir}/{protein_name}.pdb") |
| shutil.copyfile(f"{args.input_dir}/DATASET/{protein_name}/{protein_name}.csv", f"{out_protein_dir}/{protein_name}.csv") |
| |
