def main(args):
if args.cuda and not torch.cuda.is_available():
raise Exception("No GPU found, please run without --cuda")
model = UNet(n_channels=args.colordim, n_classes=args.num_class)
model2 = UNet(n_channels=args.colordim2, n_classes=args.num_class2)
if args.cuda:
model = model.cuda()
model2 = model2.cuda()
model.load_state_dict(torch.load(args.pretrain_net))
model2.load_state_dict(torch.load(args.pretrain_net2))
model.eval()
model2.eval()
predDataset = generateDataset(args.pre_root_dir,
args.img_size,
args.colordim,
isTrain=False)
predLoader = DataLoader(dataset=predDataset,
batch_size=args.predictbatchsize,
num_workers=args.threads)
with torch.no_grad():
cm_w = np.zeros((2, 2))
for batch_idx, (batch_x, batch_name) in enumerate(predLoader):
batch_x = batch_x
if args.cuda:
batch_x = batch_x.float().cuda()
out1 = model(batch_x)
prediction2 = torch.cat((batch_x, out1), 1)
out = model2(prediction2)
pred_prop, pred_label = torch.max(out, 1)
pred_label_np = pred_label.cpu().numpy()
for id in range(len(batch_name)):
predLabel_filename = args.preDir + '/' + batch_name[id] + '.png'
pred_label_single = pred_label_np[id, :, :]
label_filename = args.label_root_dir + batch_name[id] + '.png'
label = io.imread(label_filename)
cm = confusion_matrix(label.ravel(), pred_label_single.ravel())
pred_label_single = np.where(pred_label_single > 0, 255, 0)
print(np.max(pred_label_single))
print(batch_name[id])
if (np.max(pred_label_single) > 0):
io.imsave(predLabel_filename,
pred_label_single.astype(np.uint8))
#else:
#io.imsave(predLabel_filename, pred_label_single.astype(np.int32))
cm_w = cm_w + cm
#OA_s, F1_s, IoU_s = evaluate(cm)
#print('OA_s = ' + str(OA_s) + ', F1_s = ' + str(F1_s) + ', IoU = ' + str(IoU_s))
print(cm_w)
OA_w, F1_w, IoU_w = evaluate(cm_w)
print('OA_w = ' + str(OA_w) + ', F1_w = ' + str(F1_w) + ', IoU = ' +
str(IoU_w))
writer = SummaryWriter()
image = Image.open('./ht2-c2.jpg')
out = TF.to_tensor(image)
out = out.reshape(1, 3, 640, 640)
inp = torch.rand(1, 3, 640, 640)
fig = plt.figure()
plt.imshow(out[0].permute(1, 2, 0).numpy())
# plt.show
writer.add_figure("Ground Truth", fig)
fig = plt.figure()
plt.imshow(inp[0].permute(1, 2, 0).numpy())
writer.add_figure("Input", fig)
num_iter = 500
writer.add_scalar("Number_of_Iterations", num_iter)
model = UNet(3, 3)
if torch.cuda.is_available():
model.cuda()
criterion = nn.MSELoss()
learning_rate = 0.1
writer.add_scalar("Learning_Rate", learning_rate)
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)
train(num_iter, inp, out, model, optimizer, criterion)
writer.close()
in_size = 320,
train_in_size = input_size)
train_dataloader = torch.utils.data.DataLoader(SRRFDATASET, batch_size=batch_size, shuffle=True, pin_memory=True) # better than for loop
SRRFDATASET2 = ReconsDataset(all_data_path="/media/star/LuhongJin/UNC_data/SRRF/New_training_20190829/0NPY_Dataset/Dataset/Microtubule/",
maximum_intensity_4normalization_path="/home/star/0_code_lhj/DL-SIM-github/Training_codes/UNetMax_intensity.npy",
transform = ToTensor(),
training_dataset = False,
in_size = 320,
train_in_size = input_size)
validation_dataloader = torch.utils.data.DataLoader(SRRFDATASET2, batch_size=batch_size, shuffle=True, pin_memory=True) # better than for loop
model = UNet(n_channels=input_size, n_classes=output_size)
print("{} paramerters in total".format(sum(x.numel() for x in model.parameters())))
model.cuda(cuda)
optimizer = torch.optim.Adam(model.parameters(),lr=learning_rate, betas=(0.9, 0.999))
loss_all = np.zeros((2000, 4))
for epoch in range(2000):
mae_m, mae_s = val_during_training(train_dataloader)
loss_all[epoch,0] = mae_m
loss_all[epoch,1] = mae_s
mae_m, mae_s = val_during_training(validation_dataloader)
loss_all[epoch,2] = mae_m
loss_all[epoch,3] = mae_s
file = Workbook(encoding = 'utf-8')
table = file.add_sheet('loss_all')
for i,p in enumerate(loss_all):
#train_loader = DataLoader(dataset=train_set, num_workers=4, batch_size=64, shuffle=True)
#val_loader = DataLoader(dataset=val_set, num_workers=4, batch_size=1, shuffle=False)
netG = UNet(n_channels=15, n_classes=1)
#print(summary(netG,(15,128,128)))
print('# generator parameters:',
sum(param.numel() for param in netG.parameters()))
netD = Discriminator()
print('# discriminator parameters:',
sum(param.numel() for param in netD.parameters()))
#print(summary(netD,(1,256,256)))
generator_criterion = GeneratorLoss()
if torch.cuda.is_available():
netG.cuda()
netD.cuda()
generator_criterion.cuda()
optimizerG = optim.Adam(netG.parameters())
optimizerD = optim.Adam(netD.parameters())
results = {
'd_loss': [],
'g_loss': [],
'd_score': [],
'g_score': [],
'psnr': [],
'ssim': []
}
data, target = get_train_data(X_test, y_test, batch_size)
LE_img = imgread(os.path.join(dir_path_LE, "LE_01.tif"))
LE_512 = cropImage(LE_img, IMG_SHAPE[0],IMG_SHAPE[1])
sample_le = {}
for le_512 in LE_512:
tiles = crop_prepare(le_512, CROP_STEP, IMG_SIZE)
for n,img in enumerate(tiles):
if n not in sample_le:
sample_le[n] = []
img = transform.resize(img,(IMG_SIZE*2, IMG_SIZE*2),preserve_range=True,order=3)
sample_le[n].append(img)
SNR_model = UNet(n_channels=15, n_classes=15)
print("{} paramerters in total".format(sum(x.numel() for x in SNR_model.parameters())))
SNR_model.cuda(cuda)
SNR_model.load_state_dict(torch.load(SNR_model_path))
# SNR_model.load_state_dict(torch.load(os.path.join(dir_path,"model","LE_HE_mito","LE_HE_0825.pkl")))
SNR_model.eval()
SIM_UNET = UNet(n_channels=15, n_classes=1)
print("{} paramerters in total".format(sum(x.numel() for x in SIM_UNET.parameters())))
SIM_UNET.cuda(cuda)
SIM_UNET.load_state_dict(torch.load(SIM_UNET_model_path))
# SIM_UNET.load_state_dict(torch.load(os.path.join(dir_path,"model","HE_HER_mito","HE_X2_HER_0825.pkl")))
SIM_UNET.eval()
SRRFDATASET = ReconsDataset(
img_dict=sample_le,
transform=ToTensor(),
in_norm = LE_in_norm,
#quit after interrupt
import sys
dir = 'data'
ids = []
for f in os.listdir(dir + '/train'):
id = f[:-4]
ids.append([id, 0])
ids.append([id, 1])
np.random.shuffle(ids)
#%%
net = UNet(3, 1)
net.cuda()
def train(net):
optimizer = optim.Adam(net.parameters(), lr=1)
criterion = DiceLoss()
epochs = 5
for epoch in range(epochs):
print('epoch {}/{}...'.format(epoch + 1, epochs))
l = 0
for i, c in enumerate(ids):
id = c[0]
pos = c[1]
im = PIL.Image.open(dir + '/train/' + id + '.jpg')
def test(experiment_path, test_epoch):
# ========= CONFIG FILE TO READ FROM =======
config = configparser.RawConfigParser()
config.read('./' + experiment_path + '/' + experiment_path + '_config.txt')
# ===========================================
# run the training on invariant or local
path_data = config.get('data paths', 'path_local')
model = config.get('training settings', 'model')
# original test images (for FOV selection)
DRIVE_test_imgs_original = path_data + config.get('data paths', 'test_imgs_original')
test_imgs_orig = load_hdf5(DRIVE_test_imgs_original)
full_img_height = test_imgs_orig.shape[2]
full_img_width = test_imgs_orig.shape[3]
# the border masks provided by the DRIVE
DRIVE_test_border_masks = path_data + config.get('data paths', 'test_border_masks')
test_border_masks = load_hdf5(DRIVE_test_border_masks)
# dimension of the patches
patch_height = int(config.get('data attributes', 'patch_height'))
patch_width = int(config.get('data attributes', 'patch_width'))
# the stride in case output with average
stride_height = int(config.get('testing settings', 'stride_height'))
stride_width = int(config.get('testing settings', 'stride_width'))
assert (stride_height < patch_height and stride_width < patch_width)
# model name
name_experiment = config.get('experiment name', 'name')
path_experiment = './' + name_experiment + '/'
# N full images to be predicted
Imgs_to_test = int(config.get('testing settings', 'full_images_to_test'))
# Grouping of the predicted images
N_visual = int(config.get('testing settings', 'N_group_visual'))
# ====== average mode ===========
average_mode = config.getboolean('testing settings', 'average_mode')
#N_subimgs = int(config.get('training settings', 'N_subimgs'))
#batch_size = int(config.get('training settings', 'batch_size'))
#epoch_size = N_subimgs // (batch_size)
# #ground truth
# gtruth= path_data + config.get('data paths', 'test_groundTruth')
# img_truth= load_hdf5(gtruth)
# visualize(group_images(test_imgs_orig[0:20,:,:,:],5),'original')#.show()
# visualize(group_images(test_border_masks[0:20,:,:,:],5),'borders')#.show()
# visualize(group_images(img_truth[0:20,:,:,:],5),'gtruth')#.show()
# ============ Load the data and divide in patches
patches_imgs_test = None
new_height = None
new_width = None
masks_test = None
patches_masks_test = None
if average_mode == True:
patches_imgs_test, new_height, new_width, masks_test= get_data_testing_overlap(
DRIVE_test_imgs_original = DRIVE_test_imgs_original, #original'DRIVE_datasets_training_testing/test_hard_masks.npy'
DRIVE_test_groudTruth = path_data + config.get('data paths', 'test_groundTruth'), #masks
Imgs_to_test = int(config.get('testing settings', 'full_images_to_test')),
patch_height = patch_height,
patch_width = patch_width,
stride_height = stride_height,
stride_width = stride_width)
else:
patches_imgs_test, patches_masks_test = get_data_testing_test(
DRIVE_test_imgs_original = DRIVE_test_imgs_original, #original
DRIVE_test_groudTruth = path_data + config.get('data paths', 'test_groundTruth'), #masks
Imgs_to_test = int(config.get('testing settings', 'full_images_to_test')),
patch_height = patch_height,
patch_width = patch_width
)
#np.save(path_experiment + 'test_patches.npy', patches_imgs_test)
#visualize(group_images(patches_imgs_test,100),'./'+name_experiment+'/'+"test_patches")
# ================ Run the prediction of the patches ==================================
best_last = config.get('testing settings', 'best_last')
# Load the saved model
if model == 'UNet':
net = UNet(n_channels=1, n_classes=2)
elif model == 'UNet_cat':
net = UNet_cat(n_channels=1, n_classes=2)
else:
net = UNet_level4_our(n_channels=1, n_classes=2)
# load data
test_data = data.TensorDataset(torch.tensor(patches_imgs_test),torch.zeros(patches_imgs_test.shape[0]))
test_loader = data.DataLoader(test_data, batch_size=1, pin_memory=True, shuffle=False)
trained_model = path_experiment + 'DRIVE_' + str(test_epoch) + 'epoch.pth'
print(trained_model)
# trained_model= path_experiment+'DRIVE_unet2_B'+str(60*epoch_size)+'.pth'
net.load_state_dict(torch.load(trained_model))
net.eval()
print('Finished loading model :' + trained_model)
net = net.cuda()
cudnn.benchmark = True
# Calculate the predictions
predictions_out = np.empty((patches_imgs_test.shape[0],patch_height*patch_width,2))
for i_batch, (images, targets) in enumerate(test_loader):
images = Variable(images.float().cuda())
out1= net(images)
pred = out1.permute(0,2,3,1)
pred = F.softmax(pred, dim=-1)
pred = pred.data.view(-1,patch_height*patch_width,2)
predictions_out[i_batch] = pred
# ===== Convert the prediction arrays in corresponding images
pred_patches_out = pred_to_imgs(predictions_out, patch_height, patch_width, "original")
#np.save(path_experiment + 'pred_patches_' + str(test_epoch) + "_epoch" + '.npy', pred_patches_out)
#visualize(group_images(pred_patches_out,100),'./'+name_experiment+'/'+"pred_patches")
#========== Elaborate and visualize the predicted images ====================
pred_imgs_out = None
orig_imgs = None
gtruth_masks = None
if average_mode == True:
pred_imgs_out = recompone_overlap(pred_patches_out,new_height,new_width, stride_height, stride_width)
orig_imgs = my_PreProc(test_imgs_orig[0:pred_imgs_out.shape[0],:,:,:]) #originals
gtruth_masks = masks_test #ground truth masks
else:
pred_imgs_out = recompone(pred_patches_out,10,9) # predictions
orig_imgs = recompone(patches_imgs_test,10,9) # originals
gtruth_masks = recompone(patches_masks_test,10,9) #masks
# apply the DRIVE masks on the repdictions #set everything outside the FOV to zero!!
# DRIVE MASK #only for visualization
kill_border(pred_imgs_out, test_border_masks)
# back to original dimensions
orig_imgs = orig_imgs[:,:,0:full_img_height,0:full_img_width]
pred_imgs_out = pred_imgs_out[:, :, 0:full_img_height, 0:full_img_width]
gtruth_masks = gtruth_masks[:, :, 0:full_img_height, 0:full_img_width]
print ("Orig imgs shape: "+str(orig_imgs.shape))
print("pred imgs shape: " + str(pred_imgs_out.shape))
print("Gtruth imgs shape: " + str(gtruth_masks.shape))
np.save(path_experiment + 'pred_img_' + str(test_epoch) + "_epoch" + '.npy',pred_imgs_out)
# visualize(group_images(orig_imgs,N_visual),path_experiment+"all_originals")#.show()
if average_mode == True:
visualize(group_images(pred_imgs_out, N_visual),
path_experiment + "all_predictions_" + str(test_epoch) + "thresh_epoch")
else:
visualize(group_images(pred_imgs_out, N_visual),
path_experiment + "all_predictions_" + str(test_epoch) + "epoch_no_average")
visualize(group_images(gtruth_masks, N_visual), path_experiment + "all_groundTruths")
# visualize results comparing mask and prediction:
# assert (orig_imgs.shape[0] == pred_imgs_out.shape[0] and orig_imgs.shape[0] == gtruth_masks.shape[0])
# N_predicted = orig_imgs.shape[0]
# group = N_visual
# assert (N_predicted%group == 0)
# ====== Evaluate the results
print("\n\n======== Evaluate the results =======================")
# predictions only inside the FOV
y_scores, y_true = pred_only_FOV(pred_imgs_out, gtruth_masks, test_border_masks) # returns data only inside the FOV
'''
print("Calculating results only inside the FOV:")
print("y scores pixels: " + str(
y_scores.shape[0]) + " (radius 270: 270*270*3.14==228906), including background around retina: " + str(
pred_imgs_out.shape[0] * pred_imgs_out.shape[2] * pred_imgs_out.shape[3]) + " (584*565==329960)")
print("y true pixels: " + str(
y_true.shape[0]) + " (radius 270: 270*270*3.14==228906), including background around retina: " + str(
gtruth_masks.shape[2] * gtruth_masks.shape[3] * gtruth_masks.shape[0]) + " (584*565==329960)")
'''
# Area under the ROC curve
fpr, tpr, thresholds = roc_curve((y_true), y_scores)
AUC_ROC = roc_auc_score(y_true, y_scores)
# test_integral = np.trapz(tpr,fpr) #trapz is numpy integration
print("\nArea under the ROC curve: " + str(AUC_ROC))
rOc_curve = plt.figure()
plt.plot(fpr, tpr, '-', label='Area Under the Curve (AUC = %0.4f)' % AUC_ROC)
plt.title('ROC curve')
plt.xlabel("FPR (False Positive Rate)")
plt.ylabel("TPR (True Positive Rate)")
plt.legend(loc="lower right")
plt.savefig(path_experiment + "ROC.png")
# Precision-recall curve
precision, recall, thresholds = precision_recall_curve(y_true, y_scores)
precision = np.fliplr([precision])[0] # so the array is increasing (you won't get negative AUC)
recall = np.fliplr([recall])[0] # so the array is increasing (you won't get negative AUC)
AUC_prec_rec = np.trapz(precision, recall)
print("\nArea under Precision-Recall curve: " + str(AUC_prec_rec))
prec_rec_curve = plt.figure()
plt.plot(recall, precision, '-', label='Area Under the Curve (AUC = %0.4f)' % AUC_prec_rec)
plt.title('Precision - Recall curve')
plt.xlabel("Recall")
plt.ylabel("Precision")
plt.legend(loc="lower right")
plt.savefig(path_experiment + "Precision_recall.png")
# Confusion matrix
threshold_confusion = 0.5
print("\nConfusion matrix: Custom threshold (for positive) of " + str(threshold_confusion))
y_pred = np.empty((y_scores.shape[0]))
for i in range(y_scores.shape[0]):
if y_scores[i] >= threshold_confusion:
y_pred[i] = 1
else:
y_pred[i] = 0
confusion = confusion_matrix(y_true, y_pred)
print(confusion)
accuracy = 0
if float(np.sum(confusion)) != 0:
accuracy = float(confusion[0, 0] + confusion[1, 1]) / float(np.sum(confusion))
print("Global Accuracy: " + str(accuracy))
specificity = 0
if float(confusion[0, 0] + confusion[0, 1]) != 0:
specificity = float(confusion[0, 0]) / float(confusion[0, 0] + confusion[0, 1])
print("Specificity: " + str(specificity))
sensitivity = 0
if float(confusion[1, 1] + confusion[1, 0]) != 0:
sensitivity = float(confusion[1, 1]) / float(confusion[1, 1] + confusion[1, 0])
print("Sensitivity: " + str(sensitivity))
precision = 0
if float(confusion[1, 1] + confusion[0, 1]) != 0:
precision = float(confusion[1, 1]) / float(confusion[1, 1] + confusion[0, 1])
print("Precision: " + str(precision))
# Jaccard similarity index
jaccard_index = jaccard_similarity_score(y_true, y_pred, normalize=True)
print("\nJaccard similarity score: " + str(jaccard_index))
# F1 score
F1_score = f1_score(y_true, y_pred, labels=None, average='binary', sample_weight=None)
print("\nF1 score (F-measure): " + str(F1_score))
####evaluate the thin vessels
thin_3pixel_recall_indivi = []
thin_3pixel_auc_roc = []
for j in range(pred_imgs_out.shape[0]):
thick3=opening(gtruth_masks[j, 0, :, :], square(3))
thin_gt = gtruth_masks[j, 0, :, :] - thick3
thin_pred=pred_imgs_out[j, 0, :, :]
thin_pred[thick3==1]=0
thin_3pixel_recall_indivi.append(round(thin_recall(thin_gt, pred_imgs_out[j, 0, :, :], thresh=0.5), 4))
thin_3pixel_auc_roc.append(round(roc_auc_score(thin_gt.flatten(), thin_pred.flatten()), 4))
thin_2pixel_recall_indivi = []
thin_2pixel_auc_roc = []
for j in range(pred_imgs_out.shape[0]):
thick=opening(gtruth_masks[j, 0, :, :], square(2))
thin_gt = gtruth_masks[j, 0, :, :] - thick
#thin_gt_only=thin_gt[thin_gt==1]
#print(thin_gt_only)
thin_pred=pred_imgs_out[j, 0, :, :]
#thin_pred=thin_pred[thin_gt==1]
thin_pred[thick==1]=0
thin_2pixel_recall_indivi.append(round(thin_recall(thin_gt, pred_imgs_out[j, 0, :, :], thresh=0.5), 4))
thin_2pixel_auc_roc.append(round(roc_auc_score(thin_gt.flatten(), thin_pred.flatten()), 4))
#print("thin 2vessel recall:", thin_2pixel_recall_indivi)
#print('thin 2vessel auc score', thin_2pixel_auc_roc)
# Save the results
with open(path_experiment + 'test_performances_all_epochs.txt', mode='a') as f:
f.write("\n\n" + path_experiment + " test epoch:" + str(test_epoch)
+ '\naverage mode is:' + str(average_mode)
+ "\nArea under the ROC curve: %.4f" % (AUC_ROC)
+ "\nArea under Precision-Recall curve: %.4f" % (AUC_prec_rec)
+ "\nJaccard similarity score: %.4f" % (jaccard_index)
+ "\nF1 score (F-measure): %.4f" % (F1_score)
+ "\nConfusion matrix:"
+ str(confusion)
+ "\nACCURACY: %.4f" % (accuracy)
+ "\nSENSITIVITY: %.4f" % (sensitivity)
+ "\nSPECIFICITY: %.4f" % (specificity)
+ "\nPRECISION: %.4f" % (precision)
+ "\nthin 2vessels recall indivi:\n" + str(thin_2pixel_recall_indivi)
+ "\nthin 2vessels recall mean:%.4f" % (np.mean(thin_2pixel_recall_indivi))
+ "\nthin 2vessels auc indivi:\n" + str(thin_2pixel_auc_roc)
+ "\nthin 2vessels auc score mean:%.4f" % (np.mean(thin_2pixel_auc_roc))
+ "\nthin 3vessels recall indivi:\n" + str(thin_3pixel_recall_indivi)
+ "\nthin 3vessels recall mean:%.4f" % (np.mean(thin_3pixel_recall_indivi))
+ "\nthin 3vessels auc indivi:\n" + str(thin_3pixel_auc_roc)
+ "\nthin 3vessels auc score mean:%.4f" % (np.mean(thin_3pixel_auc_roc))
)
LE_512 = cropImage(LE_img, IMG_SHAPE[0], IMG_SHAPE[1])
sample_le = {}
for le_512 in LE_512:
tiles = crop_prepare(le_512, CROP_STEP, IMG_SIZE)
for n, img in enumerate(tiles):
if n not in sample_le:
sample_le[n] = []
img = transform.resize(img, (IMG_SIZE * 2, IMG_SIZE * 2),
preserve_range=True,
order=3)
sample_le[n].append(img)
SC_UNET = UNet(n_channels=15, n_classes=1)
print("{} paramerters in total".format(
sum(x.numel() for x in SC_UNET.parameters())))
SC_UNET.cuda(cuda)
SC_UNET.load_state_dict(torch.load(model_path))
# SC_UNET.load_state_dict(torch.load(os.path.join(dir_path,"model","HE_HER_mito","HE_X2_HER_0825.pkl")))
SC_UNET.eval()
SRRFDATASET = ReconsDataset(img_dict=sample_he,
transform=ToTensor(),
in_norm=LE_in_norm,
img_type=".tif",
in_size=256)
test_dataloader = torch.utils.data.DataLoader(
SRRFDATASET, batch_size=1, shuffle=False,
pin_memory=True) # better than for loop
result = np.zeros((256, 256, len(SRRFDATASET)))
for batch_idx, items in enumerate(test_dataloader):
image = items['image_in']
