import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from tensorflow.python.client import device_lib

print(device_lib.list_local_devices())

from Vnet.model_vnet3d import Vnet3dModule
from dataprocess.utils import calcu_dice
import numpy as np
import pandas as pd


def predict():
    '''
    Preprocessing for dataset
    '''
    # Read  data set (Train data from CSV file)
    csvdata = pd.read_csv('dataprocess\\data/test.csv')
    maskdata = csvdata.iloc[:, 1].values
    imagedata = csvdata.iloc[:, 0].values

    dice_values = []
    Vnet3d = Vnet3dModule(128, 128, 64, channels=4, numclass=3, costname=("dice coefficient",), inference=True,
                          model_path="log\segmeation\VNet\model\Vnet3d.pd")
    for index in range(imagedata.shape[0]):
        image_gt = np.load(imagedata[index])
        mask_pd = Vnet3d.prediction(image_gt)
        mask_gt = np.load(maskdata[index])
        mask_gt_255 = mask_gt.copy()
        mask_gt_255[mask_gt == 0] = 0
        mask_gt_255[mask_gt != 0] = 255
        dice_value = calcu_dice(mask_pd, mask_gt_255)
        print("index,dice:", (index, dice_value))
        dice_values.append(dice_value)
    average = sum(dice_values) / len(dice_values)
    print("average dice:", average)


predict()