--- a
+++ b/evaluate.py
@@ -0,0 +1,124 @@
+import argparse
+import torch
+from torch.utils.data import DataLoader
+import pandas as pd
+import os
+import numpy as np
+os.environ['KMP_DUPLICATE_LIB_OK']='True'
+
+
+from dataset import LoadDataset
+from model import InferenceNet, ECGnet
+from utils import visualize_two_PC, ECG_visual_two, visualize_PC_with_twolabel_rotated
+from loss import calculate_Dice, evaluate_pointcloud, calculate_inference_loss, calculate_reconstruction_loss
+
+def evaluate(args):
+
+    DEVICE = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
+    test_dataset = LoadDataset(path=args.partial_root, num_input=args.num_input, split='test')
+    test_dataloader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
+
+    # network = ECGnet(in_ch=args.in_ch, out_ch=args.out_ch, num_input=args.num_input, z_dims=args.z_dims)
+    network = InferenceNet(in_ch=args.in_ch, out_ch=args.out_ch, num_input=args.num_input, z_dims=args.z_dims)
+
+    network.load_state_dict(torch.load('log/net_model.pkl'))
+    network.to(DEVICE)
+
+    Dice_Scar, Dice_BZ =  [], []
+    precision_Scar, precision_BZ = [], []
+    recall_Scar, recall_BZ = [], []
+    f1_score_Scar, f1_score_BZ = [], []
+    roc_auc_Scar, roc_auc_BZ = [], []
+    pre_MI_size_Scar, pre_MI_size_BZ = [], []
+    gd_MI_size_Scar, gd_MI_size_BZ = [], []
+    MI_center_dist = []
+    MI_type_list = []
+    AHA_loc_score_list = []
+    recon_geo_list, recon_ECG_list = [], []
+
+    # testing: evaluate the mean loss
+    network.eval()
+    with torch.no_grad():
+        for i, data in enumerate(test_dataloader, 1):
+            partial_input, ECG_input, gt_MI, partial_input_coarse, MI_type = data
+            partial_input, ECG_input, gt_MI = partial_input.to(DEVICE), ECG_input.to(DEVICE), gt_MI.to(DEVICE)      
+            partial_input_coarse = partial_input_coarse.to(DEVICE)      
+            partial_input = partial_input.permute(0, 2, 1)
+
+            y_MI, y_coarse, y_detail, y_ECG, mu, log_var = network(partial_input[:, 0:7, :], ECG_input)
+            loss_geo, loss_signal = calculate_reconstruction_loss(y_coarse, y_detail, partial_input_coarse, partial_input, y_ECG, ECG_input)
+
+            Dice = calculate_Dice(y_MI, gt_MI, num_classes=3)
+            precision, recall, f1_score, roc_auc, MI_size_pre, MI_size_gd, center_distance, AHA_loc_score = evaluate_pointcloud(y_MI, gt_MI, partial_input)
+
+            Dice_Scar.append(Dice[1].cpu().detach().numpy())
+            Dice_BZ.append(Dice[2].cpu().detach().numpy())
+            precision_Scar.append(precision[1])
+            precision_BZ.append(precision[2])
+            recall_Scar.append(recall[1])
+            recall_BZ.append(recall[2])
+            # f1_score_Scar.append(f1_score[1])
+            # f1_score_BZ.append(f1_score[2])
+            # roc_auc_Scar.append(roc_auc[1])
+            # roc_auc_BZ.append(roc_auc[2])
+
+            pre_MI_size_Scar.append(MI_size_pre[1])
+            pre_MI_size_BZ.append(MI_size_pre[2])
+            gd_MI_size_Scar.append(MI_size_gd[1])
+            gd_MI_size_BZ.append(MI_size_gd[2])
+            MI_center_dist.append(center_distance)
+            AHA_loc_score_list.append(AHA_loc_score)
+            recon_geo_list.append(loss_geo.cpu().detach().numpy())
+            recon_ECG_list.append(loss_signal.cpu().detach().numpy())
+
+            MI_type_list.append(MI_type[0])
+
+            visual_check = False
+            if visual_check:
+                gd_ECG = ECG_input[0].cpu().detach().numpy()
+                y_ECG = y_ECG[0].cpu().detach().numpy()
+                ECG_visual_two(y_ECG, gd_ECG)
+                y_predict = y_MI[0].cpu().detach().numpy()
+                y_gd = gt_MI[0].cpu().detach().numpy()
+                x_input = partial_input[0].cpu().detach().numpy()
+                y_predict_argmax = np.argmax(y_predict, axis=0)
+                y_output = y_detail.permute(0, 2, 1)[0].cpu().detach().numpy()
+                visualize_PC_with_twolabel_rotated(x_input[0:3, 0:args.num_input].transpose(), y_predict_argmax, y_gd, filename='RNmap_gd_pre.pdf')
+                visualize_two_PC(x_input[0:3, 0:args.num_input].transpose(), y_output[0:3, 0:args.num_input].transpose(), y_gd, filename='PC_recon.pdf')
+
+        list = {'MI_type': MI_type_list, 'Dice_Scar': Dice_Scar, 'Dice_BZ': Dice_BZ, 'precision_Scar': precision_Scar, 'precision_BZ': precision_BZ, 
+        'recall_Scar': recall_Scar, 'recall_BZ': recall_BZ, 
+        'pre_MI_size_Scar': pre_MI_size_Scar, 'pre_MI_size_BZ': pre_MI_size_BZ,
+        'gd_MI_size_Scar': gd_MI_size_Scar, 'gd_MI_size_BZ': gd_MI_size_BZ
+        , 'MI_center_dist': MI_center_dist, 'AHA_loc_score': AHA_loc_score_list
+        , 'recon_geo': recon_geo_list, 'recon_ECG': recon_ECG_list}
+        
+        df = pd.DataFrame(list)
+        df.to_csv('MI_inference_results_sample4.csv', encoding='gbk', index=False)
+
+        print('Lei, well done!')       
+     
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--partial_root', type=str, default='./Big_data_inference/meta_data/UKB_clinical_data/')
+    parser.add_argument('--model', type=str, default='log/net_model.pkl') #'log/net_model.pkl'
+    parser.add_argument('--in_ch', type=int, default=3+4) # coordinate dimension + label index
+    parser.add_argument('--out_ch', type=int, default=3) # scar, BZ, normal
+    parser.add_argument('--z_dims', type=int, default=16)
+    parser.add_argument('--num_input', type=int, default=1024*4)
+    parser.add_argument('--batch_size', type=int, default=1)
+    parser.add_argument('--alpha', type=float, default=0.1)
+    parser.add_argument('--beta', type=float, default=1e-2)
+    parser.add_argument('--lamda', type=float, default=1)
+    parser.add_argument('--base_lr', type=float, default=1e-5) #5e-5
+    parser.add_argument('--lr_decay_steps', type=int, default=50) 
+    parser.add_argument('--lr_decay_rate', type=float, default=0.5) 
+    parser.add_argument('--weight_decay', type=float, default=1e-6) #1e-3
+    parser.add_argument('--epochs', type=int, default=500)
+    parser.add_argument('--num_workers', type=int, default=1)
+    parser.add_argument('--log_dir', type=str, default='log')
+    args = parser.parse_args()
+
+    evaluate(args)
+