GeezerVison is showing github.com — code_base/predict_cml_mia_no_fuse.py at main · datct00/code_base

Expand file tree

predict_cml_mia_no_fuse.py

Copy path

More file actions

Latest commit

History

220 lines (203 loc) · 9.44 KB

predict_cml_mia_no_fuse.py

Copy path

File metadata and controls

220 lines (203 loc) · 9.44 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

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.forward_mia_no_fuse(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.forward_mia_no_fuse(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-CML_MIA_NO_FUSE_BraTs-seed-2-Semi-CML-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))