--- a
+++ b/Main.py
@@ -0,0 +1,80 @@
+import os
+from Segmenation import BreastSeg,BreastTumor
+from ImageProcessing import read_image,Norm_Zscore,imgnorm,save_image,labeling_seg
+import argparse
+import torch
+try:
+    torch._utils._rebuild_tensor_v2
+except AttributeError:
+    def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
+        tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
+        tensor.requires_grad = requires_grad
+        tensor._backward_hooks = backward_hooks
+        return tensor
+    torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2
+    
+    
+parser = argparse.ArgumentParser(description='Tumor Segmentation from DCE-MRI')
+parser.add_argument('--cuda', type=int, default='1', required=False, help='Run in GPU')
+parser.add_argument('--pre', type=str, default='Data/Img1_pre.nii.gz', required=False, help='Image path for pre-constrast image')
+parser.add_argument('--post', type=str, default='Data/Img1_post.nii.gz', required=False, help='Image path for post-constrast image')
+parser.add_argument('--outfolder',type=str,default='Results',required=False,help='Folder for saving results')
+opt = parser.parse_args()
+print(opt)
+
+if opt.cuda and not torch.cuda.is_available():
+    raise Exception("No GPU found, please run without --cuda 0")
+
+
+
+savepath = opt.outfolder
+
+if not os.path.exists(savepath):
+    os.makedirs(savepath)
+
+# read images and output spacing information   
+print('Reading images')      
+img_pre, img_post, img_sub, scale_subject = read_image(opt)
+
+# perform the intensity normalization 1) remover outliers 2) Z-score normalization 
+# spatial normalization will be done before segmentation in the functions of BrestSeg and BreastTumor
+print('Performing image normalization')
+img_pre  = Norm_Zscore(imgnorm(img_pre))
+img_post = Norm_Zscore(imgnorm(img_post))
+img_sub  = Norm_Zscore(imgnorm(img_sub))
+
+# read models
+# the weights of the models will be loaded in the function of BrestSeg and BreastTumor
+from Models_3D import ModelBreast,ModelTumor
+
+
+if opt.cuda:
+    model_breast = ModelBreast(1,1).cuda()
+else:
+    model_breast = ModelBreast(1,1)
+if opt.cuda:
+    model_tumor = ModelTumor(3,1).cuda()
+else:
+    model_tumor = ModelTumor(3,1)
+    
+    
+# perform image segmentation    
+print('Performing image segmentation')
+breast_mask = BreastSeg(img_pre,scale_subject,model_breast,opt)
+prob_1st, seg_2nd = BreastTumor(img_sub,img_post,breast_mask,scale_subject,model_tumor,model_tumor,opt)
+
+# perform labeling
+labeled_image = labeling_seg(seg_2nd)
+
+# save images
+print('Saving segmentations')
+save_image(breast_mask,opt,'breast_mask.nii.gz')
+save_image(prob_1st,opt,'prob_1st.nii.gz')
+save_image(labeled_image,opt,'seg_2nd.nii.gz')
+print('Please check the segmentation results in the result folder ---> "Results"')
+
+
+
+
+
+