File size: 901 Bytes
be2715b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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") |