evaluate.py

# CUDA_VISIBLE_DEVICES=1,2,4,5 
import io
import os
import argparse
from tqdm import tqdm
import numpy as np
import torch
import torch.nn.functional as F
from torchvision import utils as vutils
from torchvision.transforms import GaussianBlur

from models.GPTrack_2D import GPTrack2D
# from models.RViT_BidTag import RViT
from utils.tools import get_world_size, get_global_rank, get_local_rank, get_master_ip
from utils.SSIM_metric import SSIM
from utils.PSNR_metric import PSNR
from datasets.phhk_dataset import PHHKDataset
# from datasets.pah_dataset_test import Seg_PAHDataset
from datasets.pah_all_labelled import Seg_PAHDataset
from monai.data import DataLoader

import cv2
import wandb
import matplotlib.pyplot as plt
import matplotlib.colors
import matplotlib.cm as cm

norm = matplotlib.colors.Normalize()
Gaussian = GaussianBlur(kernel_size=(5, 5), sigma=(0.1, 2.0))

class Eval:
    def __init__(self, args):
        self.RViT = GPTrack2D(image_size = args.image_size[0],
                              patch_size = args.patch_size,
                              length = args.image_size[2],
                              depth = args.num_layers,
                              heads = args.num_heads,
                              mlp_dim = args.latent_dim,
                              dropout = 0.1,).to(args.device)
        
        # pretrain_params = torch.load('/home/jyangcu/Pulmonary_Arterial_Hypertension/results/checkpoints/checkpoint_130.pth', map_location='cpu')
        # pretrain_params = {k.replace('module.', ''): v for k, v in pretrain_params.items() if k.replace('module.', '') in self.RViT.state_dict()}
        # self.RViT.load_state_dict(pretrain_params)

        args.dataset_path = r'/home/jyangcu/Dataset/dataset_pa_iltrasound_nill_files_clean_image'
        valid_dataset = Seg_PAHDataset('/home/jyangcu/Dataset/dataset_pa_iltrasound_nill_files_clean_image/label_all_frame', view_num=['4'], length = args.image_size[-1])
        self.valid_loader  = DataLoader(valid_dataset, batch_size=1, shuffle=False, num_workers=1)
        self.SSIM_metric = SSIM(window_size = 8)
        self.PSNR_metric = PSNR()

        self.eval(args)

    def eval(self, args):
        self.RViT.eval()
        record_steps = 0
        all_psnr, all_dice, all_ssim, lv_dice, la_dice, ra_dice, rv_dice = [], [], [], [], [], [], []
        pbar = tqdm(self.valid_loader)
        for step, (vids, vids_org, masks) in enumerate(pbar):
            # if step > 0:
            #     print('Finish, Stop in Step ->', step)
            #     break
            # b, c, t, h, w = vids.size
            # input_vids = vids[args.selected_view[0]]
            input_vids = vids
            hidden = torch.zeros(args.batch_size, (args.image_size[0] // args.patch_size) ** 2, args.latent_dim).to(args.device)
            _, inf_flow_all, neg_inf_flow_all, lag_flow, neg_lag_flow, lag_register, forward_regsiter, backward_regsiter = self.RViT(input_vids.transpose(1,2).to(args.device), hidden, train=False)

            com_input_vids = input_vids.squeeze(0).transpose(1,0)[1:, ...].cpu().mul(255)
            com_lag_register = lag_register.squeeze(0).cpu().mul(255)
            com_forward_regsiter = torch.stack(forward_regsiter, dim=1).squeeze(0).cpu().mul(255)

            ssim_score = self.SSIM_metric(com_forward_regsiter, com_input_vids)
            psnr_score = self.PSNR_metric(com_forward_regsiter, com_input_vids)
            all_psnr.append(psnr_score.detach().cpu().numpy())
            all_ssim.append(ssim_score.detach().cpu().numpy())

            print("lag-reg: SSIM ---> {} , PSNR ---> {}".format(ssim_score, psnr_score))

            track_segments = []
            for idx in tqdm(range(len(inf_flow_all))):                
                # vutils.save_image(input_vids[0, :, idx, ...].add(1.0).mul(0.5), f'./results/flow_result_eval/org_img_{idx}.png')

                # lag_flow_plt = self.plot_warpgrid(lag_register[0][idx], lag_flow[0, :, idx, ...], interval=4, mark='c')
                # lag_flow_plt.savefig(f'./results/flow_result_eval/lag_flow_img_warp_{idx}.png')
                # lag_flow_plt.clf()

                inf_flow_plt = self.plot_warpgrid(vids_org[0, :, idx, ...], inf_flow_all[idx][0, ...], interval=8, mark='w', heatmap=False)
                inf_flow_plt.savefig(f'./results/flow_result_eval/inf_flow_heatmap_warp_{idx}.png')
                inf_flow_plt.clf()

                # For Contours evaluation
                # c_mask = masks[:, 0, 0, ...].transpose(1,2)
                # inf_flow_seg_plt, c_mask = self.plot_seg_warpgrid(vids_org[0, :, idx+1, ...], c_mask, masks[:, 0, idx+1, ...], lag_flow[0, :, idx, ...], mark='w')
                
                # For Masks Evaluation
                if idx < 1:
                    c_mask = masks[:, 0, idx, ...].transpose(1,2)
                else:
                    c_mask = c_mask
                inf_flow_seg_plt, c_mask = self.plot_seg_warpgrid(vids_org[0, :, idx+1, ...], c_mask, masks[:, 0, idx+1, ...], inf_flow_all[idx][0, ...], mark='w')
                track_segments.append(c_mask)
                inf_flow_seg_plt.savefig(f'./results/flow_result_eval/inf_flow_seg_warp_{idx}.png')
                inf_flow_seg_plt.clf()
            
            track_segments = torch.stack(track_segments, dim = 0)
            gt_segments = masks[:, 0, 1:,...].permute(1,0,3,2)
            pixel_acc, dice, precision, specificity, recall = self._calculate_overlap_metrics(torch.where(gt_segments > 0, 1, 0), torch.where(track_segments > 0, 1, 0))
            all_dice.append(dice.detach().cpu().numpy())
            
            for i in range(4):
                s_gt_segment = torch.where(gt_segments == i+1, 1, 0)
                s_track_segments = torch.where(track_segments == i+1, 1, 0)
                _, i_dice, _, _, _ = self._calculate_overlap_metrics(s_gt_segment, s_track_segments)
                if i == 0:
                    lv_dice.append(i_dice.detach().cpu().numpy())
                elif i == 1:
                    la_dice.append(i_dice.detach().cpu().numpy())
                elif i == 2:
                    ra_dice.append(i_dice.detach().cpu().numpy())
                elif i == 3:
                    rv_dice.append(i_dice.detach().cpu().numpy())
            
            print("Pixel Acc is : ", pixel_acc)
            print("Dice Score is : ", dice)
            print("Precision is : ", precision)
            print("Specificity is : ", specificity)
            print("Recall is : ", recall)
            # orginial_imgs = input_vids[0, :, 1:, ...].transpose(0, 1)
            # forward_imgs = torch.stack(forward_regsiter, dim=0)[:, 0, ...].detach().cpu()
            # backward_imgs = torch.stack(backward_regsiter, dim=0)[:, 0, ...].detach().cpu()
            # lag_imgs = lag_register[0].detach().cpu()
            # combine_imgs = torch.cat([orginial_imgs, forward_imgs, backward_imgs, lag_imgs], dim=0)
            # vutils.save_image(combine_imgs.add(1.0).mul(0.5), os.path.join("results/example_result_eval", f"example_{step}.jpg"), nrow=len(inf_flow_all))

        print("SSIM: Mean:{}, Std:{}".format(np.mean(all_ssim), np.std(all_ssim)))
        print("PSNR: Mean:{}, Std:{}".format(np.mean(all_psnr), np.std(all_psnr)))
        print("DICE: Mean:{}, Std:{}".format(np.mean(all_dice), np.std(all_dice)))
        print("LV_DICE: Mean:{}, Std:{}".format(np.mean(lv_dice), np.std(lv_dice)))
        print("LA_DICE: Mean:{}, Std:{}".format(np.mean(la_dice), np.std(la_dice)))
        print("RA_DICE: Mean:{}, Std:{}".format(np.mean(ra_dice), np.std(ra_dice)))
        print("RV_DICE: Mean:{}, Std:{}".format(np.mean(rv_dice), np.std(rv_dice)))

    def plot_warpgrid(self, img, warp, segment_result=None, 
                      interval=2, show_axis=False, mark='k', next_frame=None, 
                      heatmap=False, get_warp_img=False):
        """
        plots the given warpgrid
        @param warp: array, H x W x 2, the transformation
        @param interval: int, The interval between grid-lines
        @param show_axis: Bool, should axes be included?
        @return: matplotlib plot. Show with plt.show()
        """
        vectors = [torch.arange(0, s) for s in (args.image_size[0], args.image_size[1])]
        grids = torch.meshgrid(vectors)
        grid = torch.stack(grids)  # y, x, z
        grid = torch.unsqueeze(grid, 0)  # add batch
        grid = grid.type(torch.FloatTensor)

        warp = warp.unsqueeze(0).detach().cpu()
        warp_save = warp
        velocity_field = torch.sqrt(torch.mul(torch.pow(warp[:, 0], 2), torch.pow(warp[:, 1], 2)))
        velocity_field = velocity_field - torch.min(velocity_field)
        velocity_field = velocity_field / torch.max(velocity_field)

        # warp[:, 0] = torch.where(velocity_field > 0.2, warp[:, 0], 0)
        # warp[:, 1] = torch.where(velocity_field > 0.2, warp[:, 1], 0)
        warp = grid + warp
        warp_save = grid + warp_save
        shape = warp.shape[2:]

        for i in range(len(shape)):
            warp[:, i, ...] = 2 * (warp[:, i, ...] / (shape[i] - 1) - 0.5)
            warp_save[:, i, ...] = (2 * (warp_save[:, i, ...] / (shape[i] - 1) - 0.5) - 2 * (grid[:, i] / (shape[i] - 1) - 0.5)) * shape[i]

        if len(shape) == 2:
            warp = warp.permute(0, 2, 3, 1)
            warp = warp[..., [1, 0]]
        elif len(shape) == 3:
            warp = warp.permute(0, 2, 3, 4, 1)
            warp = warp[..., [2, 1, 0]]

        if len(shape) == 2:
            warp_save = warp_save.permute(0, 2, 3, 1)
            warp_save = warp_save[..., [1, 0]]
        elif len(shape) == 3:
            warp_save = warp_save.permute(0, 2, 3, 4, 1)
            warp_save = warp_save[..., [2, 1, 0]]

        warp_save = warp_save[0, ...]

        if img is not None:
            img = img.transpose(1,2)
            if next_frame is not None:
                next_img = next_frame.transpose(1,2)
            
            # Get the warpping img
            if get_warp_img:
                new_locs = torch.zeros_like(warp)
                for i in range(len(shape)):
                    new_locs[:, i, ...] = 2 * (warp[:, i, ...] / (shape[i] - 1) - 0.5)
                img_warp = torch.nn.functional.grid_sample(img.unsqueeze(0), new_locs, align_corners=True, mode='bilinear')[0]
            
            # Get the heatmap according to the velocity fields
            if heatmap:
                lengths = np.sqrt(np.square(warp_save[:, :, 0]) + np.square(warp_save[:, :, 1]))
                img_heat = lengths - torch.min(lengths)
                img_heat = img_heat / torch.max(img_heat)
                img_heat = torch.where(img_heat > 0.2, 0.2, img_heat)
                img_heat = img_heat * 4

                img_heat = torch.where(img.permute(2, 1, 0)[:,:,0]> -0.5, img_heat, 0)
                img_heat = img_heat.mul(255)

            # Get the segmentation Result
            if segment_result is not None:
                # Display : mask with grey scale, mask with color 1, mask with color 2 and mask with outside contours
                filterd_masked_bw = np.zeros((256, 256, 1))
                filterd_masked_c1 = np.zeros((256, 256, 3))
                filterd_masked_c2 = np.zeros((256, 256, 3))
                filterd_masked_expand = np.zeros((256, 256, 3))

                # Here for each segmented part
                # for part in range(1, 5):
                seg_part = torch.where(torch.nn.Sigmoid()(segment_result[:, 1, ...]) > 0.5, 1, 0).permute(1, 2, 0)

                # Find the contours from segmentation results
                seg_part_cv2 = np.where(seg_part.numpy() > 0, 255, 0)
                _, threshold = cv2.threshold(np.uint8(seg_part_cv2), 127, 255, cv2.THRESH_BINARY)
                contours, _ = cv2.findContours(threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) # countours is a python list
                
                # Find the largest countours
                if len(contours) > 0:
                    area = []
                    for k in range(len(contours)):
                        area.append(cv2.contourArea(contours[k]))
                    max_idx = np.argmax(np.array(area))

                    cv2.drawContours(filterd_masked_c1, contours, max_idx, (83, 253, 254), cv2.FILLED)
                    cv2.drawContours(filterd_masked_c2, contours, max_idx, (253, 83, 254), cv2.FILLED)
                    cv2.drawContours(filterd_masked_bw, contours, max_idx, 255, cv2.FILLED)
                    
                    filterd_masked_bw = self.dilate_mask(filterd_masked_bw, 10)
                    expanded_contours, _ = cv2.findContours(np.uint8(filterd_masked_bw), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
                    if len(expanded_contours) > 0:
                        area = []
                        for k in range(len(expanded_contours)):
                            area.append(cv2.contourArea(expanded_contours[k]))
                        max_idx_expanded_contours = np.argmax(np.array(area))
                        cv2.drawContours(filterd_masked_expand, expanded_contours, max_idx_expanded_contours, (83, 253, 254), cv2.FILLED)

                        # cnt_max = expanded_contours[max_idx_expanded_contours]
                        # rect = cv2.minAreaRect(cnt_max)
                        # box = cv2.boxPoints(rect)
                        # box = [np.int0(box)]
                        # # left top / left bottom / right bottom / right top : formulated as Y, X
                        # box[0][0] = [box[0][0][0]+10, box[0][0][1]-10]
                        # box[0][1] = [box[0][1][0]+10, box[0][1][1]-10]
                        # box[0][2] = [box[0][2][0]+10, box[0][2][1]+10]
                        # box[0][3] = [box[0][3][0]+10, box[0][3][1]+10]
                        # cv2.drawContours(filterd_masked_expand_square, box, -1, (255, 255, 255), cv2.FILLED)

                # img_ref = img.permute(2, 1, 0)
                # warp_save[:,:,0] = torch.where(img_ref[:,:,0]<-0.998, 0, warp_save[:,:,0])
                # warp_save[:,:,1] = torch.where(img_ref[:,:,0]<-0.998, 0, warp_save[:,:,1])
                # velocity_lengths = torch.sqrt(torch.square(warp_save[:, :, 0]) + torch.square(warp_save[:, :, 1])).transpose(0, 1).unsqueeze(-1).repeat(1, 1, 3)
                
                filterd_masked_expand = torch.from_numpy(filterd_masked_expand)
                filterd_masked_all = (filterd_masked_expand-filterd_masked_c1+filterd_masked_c2).transpose(0, 1).numpy().astype(np.uint8)

            if heatmap:
                plt.imshow(img.permute(1, 2, 0).mul(255).detach().cpu().numpy(), cmap='gray', vmin=0, vmax=255)
                plt.imshow(img_heat.unsqueeze(0).permute(2, 1, 0).detach().cpu().numpy(), cmap='jet', vmin=0, vmax=255, alpha=0.4)
            else:
                # img = torch.flip(img, dims=[2]) 
                plt.imshow(img.permute(1, 2, 0).mul(255).detach().cpu().numpy(), cmap='gray', vmin=0, vmax=255)
                if segment_result is not None:
                    plt.imshow(filterd_masked_all, vmin=0, vmax=255, alpha=0.4)

        if show_axis is False:
            plt.axis('off')
        ax = plt.gca()
        ax.set_aspect('equal')

        # This code is for the mesh drawing
        #
        # warp = warp[0, ...].numpy()
        # for row in range(0, warp.shape[0], interval):
        #     plt.plot(warp[row, :, 1], warp[row, :, 0], mark)
        # for col in range(0, warp.shape[1], interval):
        #     plt.plot(warp[:, col, 1], warp[:, col, 0], mark)

        warp_save = warp_save.numpy()
        if segment_result is not None:
            filterd_masked_expand = np.transpose(np.sum((filterd_masked_expand).transpose(0, 1).numpy(), axis=2))
        if img is not None:
            # img_ref = img.permute(2, 1, 0)
            # warp_save[:,:,0] = np.where(img_ref[:,:,0]<-0.99, 0, warp_save[:,:,0])
            # warp_save[:,:,1] = np.where(img_ref[:,:,0]<-0.99, 0, warp_save[:,:,1])
            if segment_result is not None:
                warp_save[:,:,0] = np.where(filterd_masked_expand[:,:] > 0, warp_save[:,:,0], 0)
                warp_save[:,:,1] = np.where(filterd_masked_expand[:,:] > 0, warp_save[:,:,1], 0)

            img_ref = img.permute(2, 1, 0)
            warp_save[:,:,0] = np.where(img_ref[:,:,0]<0.01, 0, warp_save[:,:,0])
            warp_save[:,:,1] = np.where(img_ref[:,:,0]<0.01, 0, warp_save[:,:,1])

        plt.quiver(grids[0][::4, ::4], grids[1][::4, ::4], warp_save[::4, ::4, 0], warp_save[::4, ::4, 1], units='xy', scale_units='xy', angles='xy', color='r', scale=1/2)

        return plt
    
    def plot_seg_warpgrid(self, img, mask, mask_tgt, warp, segment_result=None, 
                          interval=2, show_axis=False, mark='k', wrap_seg=True):
        """
        plots the given warpgrid
        @param warp: array, H x W x 2, the transformation
        @param interval: int, The interval between grid-lines
        @param show_axis: Bool, should axes be included?
        @return: matplotlib plot. Show with plt.show()
        """
        vectors = [torch.arange(0, s) for s in (args.image_size[0], args.image_size[1])]
        grids = torch.meshgrid(vectors)
        grid = torch.stack(grids)  # y, x, z
        grid = torch.unsqueeze(grid, 0)  # add batch
        grid = grid.type(torch.FloatTensor)

        warp = warp.unsqueeze(0).detach().cpu()
        warp = grid + warp
        warp_save = grid + warp
        shape = warp.shape[2:]

        for i in range(len(shape)):
            warp[:, i, ...] = 2 * (warp[:, i, ...] / (shape[i] - 1) - 0.5)
            warp_save[:, i, ...] = (2 * (warp[:, i, ...] / (shape[i] - 1) - 0.5) - 2 * (grid[:, i] / (shape[i] - 1) - 0.5)) * shape[i]

        if len(shape) == 2:
            warp = warp.permute(0, 2, 3, 1)
            warp = warp[..., [1, 0]]

            warp_save = warp_save.permute(0, 2, 3, 1)
            warp_save = warp_save[..., [1, 0]]

        elif len(shape) == 3:
            warp = warp.permute(0, 2, 3, 4, 1)
            warp = warp[..., [2, 1, 0]]

            warp_save = warp_save.permute(0, 2, 3, 4, 1)
            warp_save = warp_save[..., [2, 1, 0]]

        warp_save = warp_save[0, ...]

        if img is not None:
            img = img.transpose(1,2)
            mask = mask.float()
            mask_tgt = mask_tgt.transpose(1,2).float()
            # Get the warpping img
            if wrap_seg:
                seg_warp = torch.nn.functional.grid_sample(mask.unsqueeze(0), warp, align_corners=True, mode='nearest')[0]
                seg_warp_a = torch.where(seg_warp > 0, 83, 0)
                seg_warp_b = torch.where(seg_warp > 0, 253, 0)
                seg_warp_c = torch.where(seg_warp > 0, 254, 0)
                seg_warp_rgb = torch.cat([seg_warp_a, seg_warp_b, seg_warp_c], dim=0)

                seg_mask_a = torch.where(mask_tgt > 0, 253, 0)
                seg_mask_b = torch.where(mask_tgt > 0, 83, 0)
                seg_mask_c = torch.where(mask_tgt > 0, 254, 0)
                seg_mask = torch.cat([seg_mask_a, seg_mask_b, seg_mask_c], dim=0)

            # Get the segmentation Result
            if segment_result is not None:
                # Display : mask with grey scale, mask with color 1, mask with color 2 and mask with outside contours
                filterd_masked_bw = np.zeros((256, 256, 1))
                filterd_masked_c1 = np.zeros((256, 256, 3))
                filterd_masked_c2 = np.zeros((256, 256, 3))
                filterd_masked_expand = np.zeros((256, 256, 3))

                # Here for each segmented part
                # for part in range(1, 5):
                seg_part = torch.where(torch.nn.Sigmoid()(segment_result[:, 1, ...]) > 0.5, 1, 0).permute(1, 2, 0)

                # Find the contours from segmentation results
                seg_part_cv2 = np.where(seg_part.numpy() > 0, 255, 0)
                _, threshold = cv2.threshold(np.uint8(seg_part_cv2), 127, 255, cv2.THRESH_BINARY)
                contours, _ = cv2.findContours(threshold, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE) # countours is a python list
                
                # Find the largest countours
                if len(contours) > 0:
                    area = []
                    for k in range(len(contours)):
                        area.append(cv2.contourArea(contours[k]))
                    max_idx = np.argmax(np.array(area))

                    cv2.drawContours(filterd_masked_c1, contours, max_idx, (83, 253, 254), cv2.FILLED)
                    cv2.drawContours(filterd_masked_c2, contours, max_idx, (253, 83, 254), cv2.FILLED)
                    cv2.drawContours(filterd_masked_bw, contours, max_idx, 255, cv2.FILLED)
                    
                    filterd_masked_bw = self.dilate_mask(filterd_masked_bw, 10)
                    expanded_contours, _ = cv2.findContours(np.uint8(filterd_masked_bw), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
                    if len(expanded_contours) > 0:
                        area = []
                        for k in range(len(expanded_contours)):
                            area.append(cv2.contourArea(expanded_contours[k]))
                        max_idx_expanded_contours = np.argmax(np.array(area))
                        cv2.drawContours(filterd_masked_expand, expanded_contours, max_idx_expanded_contours, (83, 253, 254), cv2.FILLED)
                
                filterd_masked_expand = torch.from_numpy(filterd_masked_expand)
                filterd_masked_all = (filterd_masked_expand-filterd_masked_c1+filterd_masked_c2).transpose(0, 1).numpy().astype(np.uint8)

            img = img.expand(3,-1,-1).mul(255)
            img = torch.where(seg_warp_rgb > 0, seg_warp_rgb, img)
            img = torch.where(seg_mask > 0, seg_mask, img)
            plt.imshow(img.permute(1, 2, 0).detach().cpu().numpy().astype(np.uint8), cmap='viridis', vmin=0, vmax=255)
            # plt.imshow(seg_warp.permute(1, 2, 0).detach().cpu().numpy(), vmin=0, vmax=255, alpha=0.4)
            # plt.imshow(seg_mask.permute(1, 2, 0).detach().cpu().numpy(), vmin=0, vmax=255, alpha=0.4)

        if show_axis is False:
            plt.axis('off')
        ax = plt.gca()
        ax.set_aspect('equal')

        # This code is for the mesh drawing
        #
        # warp = warp[0, ...].numpy()
        # for row in range(0, warp.shape[0], interval):
        #     plt.plot(warp[row, :, 1], warp[row, :, 0], mark)
        # for col in range(0, warp.shape[1], interval):
        #     plt.plot(warp[:, col, 1], warp[:, col, 0], mark)

        warp_save = warp_save.numpy()
        if segment_result is not None:
            filterd_masked_expand = np.transpose(np.sum((filterd_masked_expand).transpose(0, 1).numpy(), axis=2))
        if img is not None:
            # img_ref = img.permute(2, 1, 0)
            # warp_save[:,:,0] = np.where(img_ref[:,:,0]<-0.99, 0, warp_save[:,:,0])
            # warp_save[:,:,1] = np.where(img_ref[:,:,0]<-0.99, 0, warp_save[:,:,1])
            if segment_result is not None:
                warp_save[:,:,0] = np.where(filterd_masked_expand[:,:] > 0, warp_save[:,:,0], 0)
                warp_save[:,:,1] = np.where(filterd_masked_expand[:,:] > 0, warp_save[:,:,1], 0)

            img_ref = img.permute(2, 1, 0)
            warp_save[:,:,0] = np.where(img_ref[:,:,0]<-0.99, 0, warp_save[:,:,0])
            warp_save[:,:,1] = np.where(img_ref[:,:,0]<-0.99, 0, warp_save[:,:,1])

        # plt.quiver(grids[0][::2, ::2], grids[1][::2, ::2], warp_save[::2, ::2, 0], warp_save[::2, ::2, 1], units='xy', angles='uv', color='r', scale_units='xy', scale=1/2)

        return plt, seg_warp
    
    def dilate_mask(self, mask, kernel_size):
        kernel  = np.ones((kernel_size, kernel_size), np.uint8)
        dilated = cv2.dilate(mask, kernel, iterations=1)
        return dilated

    def draw_grid(self, img, grid_height = 6, grid_width = 6, line_width=5):
        height, width, _ = img.shape
        img = (img + 1) * 127.5
        # img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)

        for x in range(0, width-1, grid_width):
            cv2.line(img, (x, 0), (x, height), (255))
        for y in range(0, height-1, grid_height):
            cv2.line(img, (0, y), (width, y), (255))

        img = img / 127.5 - 1

        return img

    def vector_to_rgb(self, angle, absolute):
        """Get the rgb value for the given `angle` and the `absolute` value

        Parameters
        ----------
        angle : float
            The angle in radians
        absolute : float
            The absolute value of the gradient
        
        Returns
        -------
        array_like
            The rgb value as a tuple with values [0..1]
        """
        max_abs = np.max(absolute)
        # normalize angle
        angle = angle % (2 * np.pi)
        if angle < 0:
            angle += 2 * np.pi

        return matplotlib.colors.hsv_to_rgb((angle / 2 / np.pi, 
                                             absolute / max_abs, 
                                             absolute / max_abs))

    def _calculate_overlap_metrics(self, gt, pred, eps=1e-5):
        output = pred.reshape(-1, )
        target = gt.reshape(-1, ).float()

        tp = torch.sum(output * target)  # TP
        fp = torch.sum(output * (1 - target))  # FP
        fn = torch.sum((1 - output) * target)  # FN
        tn = torch.sum((1 - output) * (1 - target))  # TN

        pixel_acc = (tp + tn + eps) / (tp + tn + fp + fn + eps)
        dice = (2 * tp + eps) / (2 * tp + fp + fn + eps)
        precision = (tp + eps) / (tp + fp + eps)
        recall = (tp + eps) / (tp + fn + eps)
        specificity = (tn + eps) / (tn + fp + eps)

        return pixel_acc, dice, precision, specificity, recall

def main(rank, args):

    def wandb_init():
        wandb.init(
            project='Unsupervised Echocardiogram Segmentation',
            entity='jiewen-yang66',
            name='PHHK-Dataset-Deep-Tag-original',
            notes='Ver 1.0',
            save_code=True
        )
        wandb.config.update(args)

    try:
        args.local_rank
    except AttributeError:
            args.global_rank = rank
            args.local_rank = args.enable_GPUs_id[rank]
    else:
        if args.distributed:
            args.global_rank = rank
            args.local_rank = args.enable_GPUs_id[rank]

    if args.distributed:
        torch.cuda.set_device(int(args.local_rank))
        torch.distributed.init_process_group(backend='nccl',
                                             init_method=args.init_method,
                                             world_size=args.world_size,
                                             rank=args.global_rank,
                                             group_name='mtorch'
                                             )
        print('using GPU {}-{} for training'.format(
            int(args.global_rank), int(args.local_rank)
            ))

        if args.wandb:
            if args.local_rank == args.enable_GPUs_id[0]:
                wandb_init()

    else:
        if args.wandb:
            wandb_init()

    if torch.cuda.is_available(): 
        args.device = torch.device("cuda:{}".format(args.local_rank))
    else: 
        args.device = 'cpu'

    Eval(args)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="EchoNet")
    parser.add_argument('--latent-dim', type=int, default=64, help='Latent dimension n_z (default: 256)')
    parser.add_argument('--image-size', type=tuple, default=(256, 256, 33), help='Image height and width (default: (112, 112 ,16))')
    parser.add_argument('--image-channels', type=int, default=1, help='Number of channels of images (default: 3)')
    parser.add_argument('--patch-size', type=int, default=16, help='Patch height and width (default: 8)')
    parser.add_argument('--mask-size', type=int, default=16, help='The size of mask patch (default: 16)')
    parser.add_argument('--mask-ratio', type=float, default=0.8, help='The ratio of masking area in an image (default: 0.75)')
    parser.add_argument('--blurring', type=bool, default=True, help='Whether blur the image')
    parser.add_argument('--max_sample_rate', type=int, default=1, help='The maximum sampling rate for the video')
    parser.add_argument('--min_sample_rate', type=int, default=1, help='The minimum sampling rate for the video')
    parser.add_argument('--num-heads', type=int, default=8, help='The number of head of multiscale attention (default: 8)')
    parser.add_argument('--num-layers', type=int, default=2, help='The number of transformer layers')

    # parser.add_argument('--mask-size', type=int, default=8, help='The size of mask patch (default: 16)')
    # parser.add_argument('--mask-ratio', type=float, default=0.7, help='The ratio of masking area in an image (default: 0.75)')
    parser.add_argument('--selected-view', type=list, default=['4'], help='The selected view from dataset')
    parser.add_argument('--dataset-path', type=str, default='/home/jyangcu/Dataset/PH_HK_image', help='Path to data (default: /data)')
    parser.add_argument('--batch-size', type=int, default=1, help='Input batch size for training (default: 6)')
    
    parser.add_argument('--epochs', type=int, default=3000, help='Number of epochs to train (default: 50)')
    parser.add_argument('--learning-rate', type=float, default=5e-4, help='Learning rate (default: 0.0002)')
    parser.add_argument('--beta1', type=float, default=0.9, help='Adam beta param (default: 0.0)')
    parser.add_argument('--beta2', type=float, default=0.99, help='Adam beta param (default: 0.999)')
    parser.add_argument('--clip-grad', type=bool, default=False, help='perform gradient clipping in training (default: False)')

    parser.add_argument('--enable_GPUs_id', type=list, default=[8], help='The number and order of the enable gpus')
    parser.add_argument('--wandb', type=bool, default=False, help='Enable Wandb')

    args = parser.parse_args()

    # setting distributed configurations
    # args.world_size = 1
    args.world_size = len(args.enable_GPUs_id)
    args.init_method = f"tcp://{get_master_ip()}:{23455}"
    args.distributed = True if args.world_size > 1 else False

    # setup distributed parallel training environments
    if get_master_ip() == "127.0.0.1" and args.distributed:
        # manually launch distributed processes 
        torch.multiprocessing.spawn(main, nprocs=args.world_size, args=(args,))
    else:
        # multiple processes have been launched by openmpi
        args.local_rank = args.enable_GPUs_id[0]
        args.global_rank = args.enable_GPUs_id[0]
    
        main(args.local_rank, args)