| import argparse | |
| from datasets_split import Kvasir_split, BUID_split, FGADR_split, MMWHS_MR_Heart_split, MMWHS_CT_Heart_split | |
| def get_args(): | |
| parser = argparse.ArgumentParser(description='Train the UNet on images and target masks') | |
| parser.add_argument('--dataset_name', '-ds', metavar='DS', type=str, default="", help='Name of dataset') | |
| return parser.parse_args() | |
| if __name__ == '__main__': | |
| args = get_args() | |
| if args.dataset_name == "Kvasir": | |
| Kvasir_split.Kvasir_split() | |
| elif args.dataset_name == "BUID": | |
| BUID_split.BUID_split() | |
| elif args.dataset_name == "FGADR": | |
| FGADR_split.FGADR_split() | |
| elif args.dataset_name == "MMWHS_MR_Heart": | |
| MMWHS_MR_Heart_split.MMWHS_MR_Heart_split() | |
| elif args.dataset_name == "MMWHS_CT_Heart": | |
| MMWHS_CT_Heart_split.MMWHS_CT_Heart_split() | |
| else: | |
| print("Let's input dataset name") |