-
Notifications
You must be signed in to change notification settings - Fork 23
/
infer_image.py
139 lines (116 loc) · 4.16 KB
/
infer_image.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
91
92
93
94
95
96
97
98
99
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
import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["CUDA_VISIBLE_DEVICES"]="0";
import numpy as np
import cv2
from glob import glob
from tqdm import tqdm
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.utils import CustomObjectScope
from tensorflow.keras.metrics import MeanIoU
from m_resunet import ResUnetPlusPlus
from metrics import *
from sklearn.metrics import confusion_matrix
from sklearn.metrics import recall_score, precision_score
from crf import apply_crf
from tta import tta_model
from utils import create_dir
def read_image(x):
image = cv2.imread(x, cv2.IMREAD_COLOR)
image = np.clip(image - np.median(image)+127, 0, 255)
image = image/255.0
image = image.astype(np.float32)
image = np.expand_dims(image, axis=0)
return image
def read_mask(y):
mask = cv2.imread(y, cv2.IMREAD_GRAYSCALE)
mask = mask.astype(np.float32)
mask = mask/255.0
mask = np.expand_dims(mask, axis=-1)
return mask
def mask_to_3d(mask):
mask = np.squeeze(mask)
mask = [mask, mask, mask]
mask = np.transpose(mask, (1, 2, 0))
return mask
def get_mean_iou(y_true, y_pred):
y_pred = y_pred.flatten()
y_true = y_true.flatten()
# y_true = y_true.astype(np.int32)
# # y_pred = y_pred > 0.5
# y_pred = y_pred.astype(np.float32)
# current = confusion_matrix(y_true, y_pred, labels=[0, 1])
#
# # compute mean iou
# intersection = np.diag(current)
# ground_truth_set = current.sum(axis=1)
# predicted_set = current.sum(axis=0)
# union = ground_truth_set + predicted_set - intersection
# IoU = intersection / union.astype(np.float32)
# return np.mean(IoU)
y_pred = y_pred > 0.5
y_pred = y_pred.astype(np.int32)
y_true = y_true.astype(np.int32)
m = tf.keras.metrics.MeanIoU(num_classes=2)
m.update_state(y_true, y_pred)
r = m.result().numpy()
m.reset_states()
return r
def save_images(model, x_data, y_data):
for i, (x, y) in tqdm(enumerate(zip(x_data, y_data)), total=len(x_data)):
x = read_image(x)
y = read_mask(y)
## Prediction
y_pred_baseline = model.predict(x)[0] > 0.5
y_pred_crf = apply_crf(x[0]*255, y_pred_baseline.astype(np.float32))
y_pred_tta = tta_model(model, x[0]) > 0.5
y_pred_tta_crf = apply_crf(x[0]*255, y_pred_tta.astype(np.float32))
y_pred_crf = np.expand_dims(y_pred_crf, axis=-1)
y_pred_tta_crf = np.expand_dims(y_pred_tta_crf, axis=-1)
sep_line = np.ones((256, 10, 3)) * 255
## MeanIoU
miou_baseline = get_mean_iou(y, y_pred_baseline)
miou_crf = get_mean_iou(y, y_pred_crf)
miou_tta = get_mean_iou(y, y_pred_tta)
miou_tta_crf = get_mean_iou(y, y_pred_tta_crf)
print(miou_baseline, miou_crf, miou_crf, miou_tta_crf)
y1 = mask_to_3d(y) * 255
y2 = mask_to_3d(y_pred_baseline) * 255.0
y3 = mask_to_3d(y_pred_crf) * 255.0
y4 = mask_to_3d(y_pred_tta) * 255.0
y5 = mask_to_3d(y_pred_tta_crf) * 255.0
# y2 = cv2.putText(y2, str(miou_baseline), (0, 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
all_images = [
x[0] * 255,
sep_line, y1,
sep_line, y2,
sep_line, y3,
sep_line, y4,
sep_line, y5
]
cv2.imwrite(f"results/{i}.png", np.concatenate(all_images, axis=1))
if __name__ == "__main__":
tf.random.set_seed(42)
np.random.seed(42)
model_path = "files/resunetplusplus.h5"
create_dir("results/")
## Parameters
image_size = 256
batch_size = 32
lr = 1e-4
epochs = 5
## Validation
valid_path = "new_data/valid/"
valid_image_paths = sorted(glob(os.path.join(valid_path, "image", "*.jpg")))
valid_mask_paths = sorted(glob(os.path.join(valid_path, "mask", "*.jpg")))
with CustomObjectScope({
'dice_loss': dice_loss,
'dice_coef': dice_coef,
'bce_dice_loss': bce_dice_loss,
'focal_loss': focal_loss,
'tversky_loss': tversky_loss,
'focal_tversky': focal_tversky
}):
model = load_model(model_path)
save_images(model, valid_image_paths, valid_mask_paths)