predict_suponly.py

import os
import time
import argparse
import logging
import numpy as np
import pandas as pd
import torch
from medpy import metric
import skimage.io as io
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt
from utils.metrics import dice,dice_3
from utils.util import read_list,set_logging
from nets.unet import UNet
import SimpleITK as sitk
plt.switch_backend('agg')

def save_imgs(img_mode,out_path,img_np,true_mask,pred_mask,patientSliceID,dice_test,exp):
    img_uint8 = np.clip(img_np, 0, 1)   # nếu dữ liệu đã normalize về [0,1]
    img_uint8 = (img_uint8 * 255).astype(np.uint8)
    io.imsave('{}/{}_img{}.jpg'.format(out_path,patientSliceID,exp),img_uint8)

    '''save true mask'''
    true_mask = true_mask * 255
    mask_true_img = Image.fromarray(true_mask)
    mask_true_img.save('{}/{}_true{}.png'.format(out_path,patientSliceID,exp), cmap="gray")

    '''save pred mask'''
    pred_mask_gray = pred_mask * 255
    pred_mask_img = Image.fromarray(pred_mask_gray)
    pred_mask_img.save('{}/{}_pred_{:.4f}{}.png'.format(out_path,patientSliceID,dice_test,exp), cmap="gray")
    # np.save('{}/{}_pred_{:.4f}.npy'.format(out_path,patientSliceID,dice_test),pred_mask)# (160, 160) dtype('uint8')0 1 2 4
    def save_plt(num,np,title):
        plt.subplot(2,2,num)
        plt.imshow(np, cmap="gray")
        plt.axis('off')
        plt.title(title)

    save_plt(1,img_np,img_mode)
    save_plt(2,img_np,img_mode)
    save_plt(3,true_mask,"true mask")
    save_plt(4,pred_mask,title="pred mask-dice:{:.4f}".format(dice_test))
    # patientSliceID = patientSliceID.replace("_slice_","_")
    plt.savefig('{}/{}{}.png'.format(out_path,patientSliceID,exp))
    plt.show()

def predict_img(net,device,true_mask,init_img,scale_factor=1,out_threshold=0.5):
    """
    test one image
    """
    # step 1/4 : path --> img_chw
    if len(init_img.shape) == 2:
        img_chw = np.expand_dims(init_img, axis=0)
    # img_chw = img_chw / 255
    # step 2/4 : img --> tensor
    img_tensor = torch.tensor(img_chw).to(torch.float32)
    img_tensor.unsqueeze_(0)  #or img = img.unsqueeze(0)
    img_tensor = img_tensor.to(device)
    # step 3/4 : tensor --> features
    time_tic = time.time()
    with torch.no_grad():
        outs = net(img_tensor)
    if len(outs) != 2:
        outputs = outs
    else:
        outputs = outs[0]
        if torch.min(outputs) < 0 or torch.max(outputs) > 1:
            outputs = torch.sigmoid(outputs)
    time_toc = time.time()
    time_s = time_toc - time_tic
    # logging.info("time is ",time_s)
    pred_mask = outputs.ge(out_threshold).cpu().data.numpy().astype("uint8")
    pred_mask = pred_mask.squeeze()
    dice_test = dice(pred_mask, true_mask)
    return pred_mask,dice_test,time_s

def test_images(net,device,img_path,true_path,in_files,out_path=None,is_save_img=True,exp=''):
    img_mode = img_path.split('/')[-1].split('_')[-1]
    dice_total = 0.
    mean_dice= 0.
    for i, file_name in enumerate(in_files):
        # logging.info("{}.Predicting image {} ...".format(i,file_name))
        path_file = os.path.join(img_path, file_name+'.npy')
        true_file = os.path.join(true_path, file_name+'.npy')
        namesplit = os.path.splitext(file_name)
        patientSliceID = namesplit[0]
        # img = Image.open(fn)  ##when data is image,jpg,png...
        img = np.load(path_file)       ##when data is npy
        true_mask = np.load(true_file)
        pred_mask ,dice_test ,time_s = predict_img(net,device,true_mask,init_img=img,scale_factor=1,out_threshold=0.5)
        dice_total += dice_test 
        # logging.info("Patient [{}] dice is [{:.4f}]. Time consuming is {:.4f}".format(patientSliceID,dice_test,time_s))
        if is_save_img==True:
            save_imgs(img_mode,out_path,img,true_mask,pred_mask,patientSliceID,dice_test,exp)
    mean_dice = dice_total / (i+1)
    return mean_dice

def test_volumes(net,device,img_path,true_path,in_files,out_path=None):
    img_mode = img_path.split('/')[-1].split('_')[-1]
    slice_len = len(in_files)
    path_file1 = os.path.join(true_path, in_files[0]+'.npy')
    w,h = np.load(path_file1).shape
    test_img_tensor = np.zeros((slice_len,1,w,h),np.float32)
    image_3d_true = np.zeros((slice_len,w,h),np.uint8)
    for i, file_name in enumerate(in_files):
        # logging.info("{}.Predicting image {} ...".format(i,file_name))
        path_file = os.path.join(img_path, file_name+'.npy')
        true_file = os.path.join(true_path, file_name+'.npy')
        patient_name = file_name.split('.')[0]
        patient_slice = patient_name.split('_')[-1]
        patient_slice = int(patient_slice) - 1
        # img = Image.open(fn)  ##when data is image,jpg,png...
        img = np.load(path_file)       ##when data is npy
        true_mask = np.load(true_file)
        # step 1/4 : path --> img_chw
        if len(img.shape) == 2:
            img_chw = np.expand_dims(img, axis=0)
        test_img_tensor[patient_slice] = img_chw
        image_3d_true[patient_slice] = true_mask
    img_tensor = torch.tensor(test_img_tensor).to(device,torch.float32)
    with torch.no_grad():
        outs = net(img_tensor)
    if len(outs) != 2:
        outputs = outs
    else:
        outputs = outs[0]
        if torch.min(outputs) < 0 or torch.max(outputs) > 1:
            outputs = torch.sigmoid(outputs)
    pred_mask_tensor = outputs.ge(0.5).cpu().data.numpy().astype("uint8")
    image_3d_pred = np.squeeze(pred_mask_tensor,axis=1)
    return image_3d_pred,image_3d_true

def compute_3d_dice(net,device,img_path_full,true_path_full,in_files_full_list,patientID_list,out_dir):
    file_names = in_files_full_list      
    patient_dice_3d_t = 0.
    patient_ppv_3d_t = 0.
    patient_sen_3d_t = 0.
    patient_asd_3d_t = 0.
    patient_hd95_3d_t = 0.
    for i , id in enumerate(patientID_list):
        img_names = list(filter(lambda x: x.startswith(id), file_names))
        imgnp_3d_pred,imgnp_3d_true = test_volumes(net,device,img_path_full,true_path_full,in_files=img_names)
        patient_dice_3d, patient_ppv_3d, patient_sen_3d = dice_3(imgnp_3d_pred, imgnp_3d_true)
        patient_dice_3d_t += patient_dice_3d
        patient_ppv_3d_t += patient_ppv_3d
        patient_sen_3d_t += patient_sen_3d
        if np.sum(imgnp_3d_pred)==0:
            asd_3d = 0
            hd95_3d = 0
        else:
            asd_3d = metric.binary.asd(imgnp_3d_pred, imgnp_3d_true)
            hd95_3d = metric.binary.hd95(imgnp_3d_pred, imgnp_3d_true)
        patient_asd_3d_t += asd_3d
        patient_hd95_3d_t += hd95_3d
        
        logging.info("{} Patient {}'s 3d dice is {:.4f}".format(i+1,id,patient_dice_3d))
    mean_dice_all_3d = patient_dice_3d_t / (i+1)
    mean_ppv_all_3d = patient_ppv_3d_t / (i+1)
    mean_sen_all_3d = patient_sen_3d_t / (i+1)
    mean_asd_all_3d = patient_asd_3d_t / (i+1)
    mean_hd95_all_3d = patient_hd95_3d_t / (i+1)
    logging.info("Mean 3d asd on all patients: {:.4f} ".format(mean_asd_all_3d))
    logging.info("Mean 3d hd95 on all patients: {:.4f} ".format(mean_hd95_all_3d))
    logging.info("Mean 3d ppv on all patients: {:.4f} ".format(mean_ppv_all_3d))
    logging.info("Mean 3d sen on all patients: {:.4f} ".format(mean_sen_all_3d))
    return mean_dice_all_3d

if __name__ == "__main__":
    # logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    gpu_list = [0] #[0,1]
    gpu_list_str = ','.join(map(str, gpu_list))
    os.environ.setdefault("CUDA_VISIBLE_DEVICES", gpu_list_str)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device {device}')

    """Test master directory """
    predict_file = 'res-t2_t1ce-BraTs-seed-2-OnlySup-Baseline-10%_seed_1111/'
    base_dir = './' + predict_file
    log_path = os.path.join(base_dir, 'predict_t1ce.log') 
    set_logging(log_path=log_path)

    data_path = "A:\Dat\semi-CML-public\dataset\BraTS_slice"
    img_mode = 't1ce'
    out_path = os.path.join(base_dir, 'outputs_new') 
    if not os.path.exists(out_path):
        os.makedirs(out_path)

    model_path = os.path.join(base_dir, 'model_UNet_last_mode2.pth')  
    logging.info("Loading model {}".format(model_path))
    net = torch.load(model_path, map_location=device,weights_only=False)    
    net.to(device=device)
    net.eval()
    logging.info("Model loaded !")
    img_path = '{}/imgs_{}'.format(data_path,img_mode)
    true_path = '{}/masks'.format(data_path)

    """1. Test some images and save them"""
    test_list_path = data_path+'/randP1_slice_nidus_val.list'
    in_files_list = read_list(test_list_path)
    # logging.info(in_files)
    logging.info("data number:{}".format(len(in_files_list)))
    mean_dice = test_images(net,device,
                    img_path,true_path,in_files_list,out_path,is_save_img=False,exp='')
    logging.info("mean dice is {:.4f}".format(mean_dice))

    """2. Calculate mean 3D Dice"""
    time_tic_test = time.time()
    full_test_list_path = '{}/randP1_slice_all_val.list'.format(data_path)
    in_files_full_list = read_list(full_test_list_path)
    patientID_test_list_path = '{}/randP1_volume_val.list'.format(data_path)
    patientID_list = read_list(patientID_test_list_path)
    # logging.info(in_files_full_list)
    mean_dice_all_3d = compute_3d_dice(net,device,
                            img_path,true_path,in_files_full_list,patientID_list,out_dir=base_dir)
    time_end_test = time.time()
    logging.info("3D Dice test time :{:.2f} min".format((time_end_test-time_tic_test)/60))
    logging.info("Mean 3d dice on all patients:{:.4f} ".format(mean_dice_all_3d))