-
Notifications
You must be signed in to change notification settings - Fork 1
/
segmentation_based_model_predict.py
90 lines (65 loc) · 3.22 KB
/
segmentation_based_model_predict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
# Imports
import numpy as np
import _pickle as cPickle
import os
import time
# Segmentation Based Model Imports
from code.segmentation_based_model.keypoints_from_mask_utilities import *
# GLOBAL VARIABLES
# FOLDS
FOLDS = [i for i in range(5)]
# Plot and Save Masks during folds iteration
PLOT_AND_SAVE_MASKS = False
# ISBI Model Directories
isbi_model_results_dir = 'results/isbi-model'
# ISBI Model Predictions Directory
isbi_model_predictions_dir = os.path.join(isbi_model_results_dir, 'predictions')
if os.path.isdir(isbi_model_predictions_dir) == False:
os.mkdir(isbi_model_predictions_dir)
# Segmentation Based Model Directories
segmentation_based_model_results_dir = 'results/segmentation-based-model'
# Segmentation Based Model U-Net++ Directory
segmentation_based_model_unet_pp_dir = os.path.join(segmentation_based_model_results_dir, 'unet-pp')
# Segmentation Based Model U-Net++ Predicted Masks Directory
segmentation_based_model_unet_pp_masks_dir = os.path.join(segmentation_based_model_unet_pp_dir, 'predictions')
# Segmentation Based Model Final Predictions Directory
segmentation_based_model_predictions = os.path.join(segmentation_based_model_results_dir, 'predictions')
if os.path.isdir(segmentation_based_model_predictions) == False:
os.mkdir(segmentation_based_model_predictions)
# Iterate through folds
for fold in FOLDS:
print('Current fold: {}'.format(fold))
# Open U-Net++ Predicted Masks
with open(os.path.join(segmentation_based_model_unet_pp_masks_dir, 'unet_pp_preds_CV5_Fold_{}.pickle'.format(fold)), 'rb') as f:
masks = cPickle.load(f)
# Open ISBI Model Predicted Keypoints
with open(os.path.join(isbi_model_predictions_dir, 'isbi_preds_w_only_CV_Fold_{}.pickle'.format(fold)), 'rb') as f:
isbi = cPickle.load(f)
# Data preprocessing (masks)
masks = np.array(masks).reshape((-1, 384, 512))
masks *= 255
# If PLOT_AND_SAVE_MASKS == True, this will be activated
if PLOT_AND_SAVE_MASKS:
for i in range(masks.shape[0]):
plt.axis('off')
plt.imshow(masks[i], cmap='gray')
plt.savefig('mask'+str(i), bbox_inches='tight', pad_inches=0.0)
plt.show()
# Start measuring time
startTime = time.time()
# Extract contours from U-Net++ predicted masks
masks_contours = get_keypoints_from_breasts_masks(masks)
# Uncomment next line if you want to plot masks and detected contours
# plot_keypoints_from_breast_masks_contours(masks, masks_contours)
# Generate Segmentation Based Model predictions
mixed = mix_isbi_and_contours(masks_contours, isbi_preds=isbi)
# Uncomment next line if you want to plot masks and final segmentation based model predictions
# plot_keypoints_from_breast_masks_contours(masks, mixed)
# Convert predictions to ISBI notation
mixed = mixed_to_our_notation(mixed)
# Elapsed time computations for algorithm execution performance
print('The script took {} second for fold {}!'.format(time.time() - startTime, fold))
# Save predictions file
with open(os.path.join(segmentation_based_model_predictions, 'mixed_preds_CV_Fold_{}.pickle'.format(fold)), 'wb') as f:
cPickle.dump(mixed, f, -1)
print('Segmentation Based Model Predictions finished.')