--- a
+++ b/tools/test_acdc_leadboard.py
@@ -0,0 +1,253 @@
+import os
+import sys
+import numpy as np
+import nibabel as nib
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import _init_paths
+from libs.network import U_Net, U_NetDF
+from utils.image_list import to_image_list
+import libs.datasets.augment as standard_augment
+import libs.datasets.joint_augment as joint_augment
+
+def progress_bar(curr_idx, max_idx, time_step, repeat_elem = "_"):
+    max_equals = 55
+    step_ms = int(time_step*1000)
+    num_equals = int(curr_idx*max_equals/float(max_idx))
+    len_reverse =len('Step:%d ms| %d/%d ['%(step_ms, curr_idx, max_idx)) + num_equals
+    sys.stdout.write("Step:%d ms|%d/%d [%s]" %(step_ms, curr_idx, max_idx, " " * max_equals,))
+    sys.stdout.flush()
+    sys.stdout.write("/b" * (max_equals+1))
+    sys.stdout.write(repeat_elem * num_equals)
+    sys.stdout.write("/b"*len_reverse)
+    sys.stdout.flush()
+    if curr_idx == max_idx:
+        print('/n')
+
+def load_nii(img_path):
+    """
+    Function to load a 'nii' or 'nii.gz' file, The function returns
+    everyting needed to save another 'nii' or 'nii.gz'
+    in the same dimensional space, i.e. the affine matrix and the header
+
+    Parameters
+    ----------
+
+    img_path: string
+    String with the path of the 'nii' or 'nii.gz' image file name.
+
+    Returns
+    -------
+    Three element, the first is a numpy array of the image values,
+    the second is the affine transformation of the image, and the
+    last one is the header of the image.
+    """
+    nimg = nib.load(img_path)
+    return nimg.get_fdata(), nimg.affine, nimg.header
+
+def save_nii(vol, affine, hdr, path, prefix, suffix):
+    vol = nib.Nifti1Image(vol, affine, hdr)
+    vol.set_data_dtype(np.uint8)
+    nib.save(vol, os.path.join(path, prefix+'_'+suffix + ".nii.gz"))
+
+
+def get_person_names(root_path=None):
+    persons_name = os.listdir(root_path)
+    persons_name = [pn for pn in persons_name if "patient" in pn]
+    persons_name.sort()
+    return persons_name
+
+def get_patient_data(patient, root_path):
+    patient_data = {}
+    infocfg_p = os.path.join(root_path, patient, "Info.cfg")
+
+    with open(infocfg_p) as f_in:
+        for line in f_in:
+            l = line.rstrip().split(": ")
+            patient_data[l[0]] = l[1]
+
+    ed_path = os.path.join(root_path, patient, "%s_frame%02d.nii.gz" % (patient, int(patient_data['ED'])))
+    es_path = os.path.join(root_path, patient, "%s_frame%02d.nii.gz" % (patient, int(patient_data['ES'])))
+    img_4d_path = os.path.join(root_path, patient, "{}_4d.nii.gz".format(patient))
+    # ed_gt_path = os.path.join(root_path, patient, "%s_frame%02d_gt.nii.gz" % (patient, int(patient_data['ED'])))
+    # es_gt_path = os.path.join(root_path, patient, "%s_frame%02d_gt.nii.gz" % (patient, int(patient_data['ES'])))
+
+    ed, affine, hdr = load_nii(ed_path)
+    patient_data['ED_VOL'] = np.swapaxes(ed, 0, -1)
+    patient_data['3D_affine'] = affine
+    patient_data['3D_hdr'] = hdr
+
+    es, _, _ = load_nii(es_path)  # (w, h, slices)
+    patient_data['ES_VOL'] = np.swapaxes(es, 0, -1)
+
+    img_4d, affine_4d, hdr_4d = load_nii(img_4d_path)  # (w, h, slices, times)
+    patient_data['4D'] = np.swapaxes(img_4d, 0, 1)
+    patient_data['4D_affine'] = affine_4d
+    patient_data['4D_hdr'] = hdr_4d
+
+    patient_data['size'] = img_4d.shape[:2][::-1]
+    patient_data['pid'] = patient
+
+    # ed_gt = load_nii(ed_gt_path)
+    # patient_data['ED_GT'] = np.swapaxes(ed_gt, 0, 1)
+
+    # es_gt = load_nii(es_gt_path)
+    # patient_data['ES_GT'] = np.swapaxes(es_gt, 0, 1)
+    return patient_data
+
+def test_it(model, data, device="cuda", used_df=True):
+    model.eval()
+    imgs = data
+
+    imgs = imgs.to(device)
+    # gts = gts.to(device)
+
+    net_out = model(imgs)
+    if used_df:
+        preds_out = net_out[0]
+        preds_df = net_out[1]
+    else:
+        preds_out = net_out[0]
+        preds_df = None
+    preds_out = nn.functional.softmax(preds_out, dim=1)
+    _, preds = torch.max(preds_out, 1)
+    preds = preds.unsqueeze(1)  # (N, 1, *)
+
+    return preds, preds_df
+
+def transform(imgs):
+    mean = 63.19523533061758
+    std = 70.74166957523165
+    trans = standard_augment.Compose([standard_augment.To_PIL_Image(),
+                                    # joint_augment.RandomAffine(0,translate=(0.125, 0.125)),
+                                    # joint_augment.RandomRotate((-180,180)),
+                                    # joint_augment.FixResize(224),
+                                    standard_augment.to_Tensor(),
+                                    standard_augment.normalize([mean], [std]),
+                                      ])
+    return trans(imgs)
+
+def test_voxel(model, imgs, used_df, multi_batches=False, resize=None):
+    """ imgs: (slices, H, W)
+        preds: (slices, 1, H, W)
+    """
+    imgs = imgs[..., None].astype(np.float32)
+    B, _, _, C = imgs.shape
+
+    if multi_batches:
+        data, origin_shape = to_image_list(imgs, size_divisible=32, return_size=True)
+        preds, _ = test_it(model, data)
+        
+        # for j in range(imgs.shape[0]):
+        #     preds[j, ...] = pred.cpu().numpy()[j, :, :origin_shape[j][0], :origin_shape[j][1]]
+    else:
+        preds = torch.zeros(B, C, resize[0], resize[1])
+        for j, pt in enumerate(imgs):
+            data = [transform(pt)]
+            data, origin_shape = to_image_list(data, size_divisible=32, return_size=True)
+            pred, _ = test_it(model, data, used_df=used_df)
+            preds[j, ...] = pred[0, 0, :origin_shape[0][0], :origin_shape[0][1]]
+
+    if resize is not None:
+        # preds = F.interpolate(preds, size=resize, mode='nearest')
+        preds = preds[..., :resize[0], :resize[1]]
+
+    return preds.cpu().numpy()[:, 0, ...]
+
+def create_model(model_name, selfeat):
+    if model_name == 'U_NetDF':
+        model = U_NetDF(selfeat=selfeat, auxseg=True)
+    elif model_name == 'U_Net':
+        model = U_Net()
+    # elif model_name == 'Resnet18_DfUnet':
+    #     model = Resnet18_DfUnet()
+    # elif model_name == 'DenseNet':
+    #     model = DenseNet()
+    # elif model_name == 'DenseNet_DF':
+    #     model = DenseNet_DF(selfeat=selfeat)
+
+    return model
+
+def test(mgpus, model_name, model_path, used_df, selfeat, log_path):
+
+    model = create_model(model_name, selfeat=selfeat)
+    if mgpus is not None and len(mgpus) > 2:
+        model = nn.DataParallel(model)
+    model.cuda()
+
+    checkpoint = torch.load(model_path, map_location='cpu')
+    model.load_state_dict(checkpoint['model_state'])
+
+    root_path = "MICCAIACDC2017/ACDC_DataSet/testing/testing/"
+    root_path = "/root/ACDC_DataSet/testing/testing/"
+    persons_name = get_person_names(root_path)
+    for j, pn in enumerate(persons_name):
+        s_time = time.time()
+        patient_data = get_patient_data(pn, root_path)
+
+        # (slices, h, w)
+        es_pred = test_voxel(model, patient_data['ES_VOL'], used_df=used_df, resize=patient_data['size'])
+        ed_pred = test_voxel(model, patient_data['ED_VOL'], used_df=used_df, resize=patient_data['size'])
+        es_pred = np.transpose(es_pred, (2, 1, 0))
+        ed_pred = np.transpose(ed_pred, (2, 1, 0))
+
+        img_4D = patient_data['4D']
+        h, w, s, t = img_4D.shape
+        pred_4D = np.zeros((w, h, s, t))
+        for i in range(img_4D.shape[-1]):
+            pred = test_voxel(model, np.transpose(img_4D[...,i], (2, 0, 1)), used_df=used_df, resize=patient_data['size'])
+            pred_4D[..., i] = np.transpose(pred, (2, 1, 0))
+        
+        save_path = os.path.join(log_path, "all_predictions")
+        os.makedirs(save_path, exist_ok=True)
+        CheckSizeAndSaveVolume(pred_4D, patient_data, save_path)
+        progress_bar(j%(len(persons_name)+1), len(persons_name),time.time()-s_time)
+
+
+def CheckSizeAndSaveVolume(seg_4D, patient_data, save_path):
+    """
+    TODO:
+    """ 
+    prefix = patient_data['pid']
+    suffix = '4D'
+
+    # save_nii(seg_4D, patient_data['4D_affine'], patient_data['4D_hdr'], save_path, prefix, suffix)
+    suffix = 'ED'
+    ED_phase_n = int(patient_data['ED'])
+    ED_pred = seg_4D[:,:,:,ED_phase_n]
+    save_nii(ED_pred.astype(np.uint8), patient_data['3D_affine'], patient_data['3D_hdr'], save_path, prefix, suffix)
+
+    suffix = 'ES'
+    ES_phase_n = int(patient_data['ES'])
+    ES_pred = seg_4D[:,:,:,ES_phase_n]
+    save_nii(ES_pred.astype(np.uint8), patient_data['3D_affine'], patient_data['3D_hdr'], save_path, prefix, suffix)
+
+    # ED_GT = patient_data.get('ED_GT', None)
+    results = []
+    return results
+
+
+
+if __name__ == "__main__":
+    # get_person_names()
+    import argparse
+    parser = argparse.ArgumentParser(description="arg parser")
+    parser.add_argument('--mgpus', type=str, default='0', required=False, help='whether to use multiple gpu')
+    parser.add_argument('--used_df', type=str, default='True', help='whether to use df')
+    parser.add_argument('--model', type=str, default='', help='whether to use df')
+    parser.add_argument('--selfeat', type=bool, default=True, help='whether to use feature select')
+    parser.add_argument('--model_path', type=str, default=None, help='whether to train with evaluation')
+    parser.add_argument('--output_dir', type=str, default="logs/acdc_logs/logs_supcat_auxseg/predtions", required=False, help='specify an output directory if needed')
+    parser.add_argument('--log_file', type=str, default="../log_predtion.txt", help="the file to write logging")
+
+    args = parser.parse_args()
+    if args.mgpus is not None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = args.mgpus
+
+    model_path = "logs/acdc_logs/logs_supcat_auxseg/ckpt/checkpoint_epoch_70.pth"
+
+    test(args.mgpus, "U_NetDF", model_path, args.used_df, args.selfeat, args.output_dir)