def test(params, self):
"""Predict masks for all images in a given directory, and save them
Args:
params (dict): the parameters of the network
"""
# Get the testing parameters
perform_watershed = params['watershed']
bbox_min_score = params['min_score']
nms_thresh = params['nms_threshold']
postProcess = params['postProcess']
resize_scale = params['scale_ratio']
# Load the data
# x_test, y_test: test images and corresponding labels
x_id, x_test = load_data_test(self.batch_seg_path)
# pred_dict and pred_dict_final save all the temp variables
pred_dict_final = {}
train_initial = tf.placeholder(dtype=tf.float32, shape=[1, None, None, 1])
input_shape = tf.shape(train_initial)
input_height = input_shape[1]
input_width = input_shape[2]
im_shape = tf.cast([input_height, input_width], tf.float32)
# number of classes needed to be classified, for our case this equals to 2
# (foreground and background)
nb_classes = 2
# feed the initial image to U-Net, we expect 2 outputs:
# 1. feat_map of shape (?,hf,wf,1024), which will be passed to the
# region proposal network
# 2. final_logits of shape(?,h,w,2), which is the prediction from U-net
with tf.variable_scope('model_U-Net') as scope:
final_logits, feat_map = UNET(nb_classes, train_initial)
# The final_logits has 2 channels for foreground/background softmax scores,
# then we get prediction with larger score for each pixel
pred_masks = tf.argmax(final_logits, axis=3)
pred_masks = tf.reshape(pred_masks, [input_height, input_width])
pred_masks = tf.to_float(pred_masks)
# Dynamic anchor base size calculated from median cell lengths
base_size = anchor_size(tf.reshape(pred_masks,
[input_height, input_width]))
# scales and ratios are used to generate different anchors
scales = np.array([0.5, 1, 2])
ratios = np.array([0.125, 0.25, 0.5, 1, 2, 4, 8])
# stride is to control how sparse we want to place anchors across the image
# stride = 16 means to place an anchor every 16 pixels on the original image
stride = 16
# Generate the anchor reference with respect to the original image
ref_anchors = generate_anchors_reference(base_size, ratios, scales)
num_ref_anchors = scales.shape[0] * ratios.shape[0]
feat_height = input_height / stride
feat_width = input_width / stride
# Generate all the anchors based on ref_anchors
all_anchors = generate_anchors(ref_anchors, stride,
[feat_height, feat_width])
num_anchors = all_anchors.shape[0]
with tf.variable_scope('model_RPN') as scope:
prediction_dict = RPN(feat_map, num_ref_anchors)
# Get the tensors from the dict
rpn_cls_prob = prediction_dict['rpn_cls_prob']
rpn_bbox_pred = prediction_dict['rpn_bbox_pred']
proposal_prediction = RPNProposal(rpn_cls_prob, rpn_bbox_pred, all_anchors,
im_shape, nms_thresh)
pred_dict_final['all_anchors'] = tf.cast(all_anchors, tf.float32)
prediction_dict['proposals'] = proposal_prediction['proposals']
prediction_dict['scores'] = proposal_prediction['scores']
pred_dict_final['rpn_prediction'] = prediction_dict
scores = pred_dict_final['rpn_prediction']['scores']
proposals = pred_dict_final['rpn_prediction']['proposals']
pred_masks_watershed = tf.to_float(
marker_watershed(scores,
proposals,
pred_masks,
min_score=bbox_min_score))
# start point for testing, and end point for graph
sess = tf.Session()
sess.run(tf.global_variables_initializer())
num_batches_test = len(x_test)
saver = tf.train.Saver()
masks1 = []
# Restore the per-image normalization model from the trained network
saver.restore(sess, './Network/whole_norm.ckpt')
sess.run(tf.local_variables_initializer())
for j in tqdm(range(0, num_batches_test)):
# whole image normalization
batch_data = x_test[j]
batch_data_shape = batch_data.shape
image = np.reshape(batch_data,
[batch_data_shape[0], batch_data_shape[1]])
if resize_scale != 1:
image = rescale(image,
self.params['scale_ratio'],
anti_aliasing=True)
# Clip the height and width to be 16-fold
imheight, imwidth = image.shape
imheight = imheight // 16 * 16
imwidth = imwidth // 16 * 16
image = image[:imheight, :imwidth]
image_normalized_wn = whole_image_norm(image)
image_normalized_wn = np.reshape(image_normalized_wn,
[1, imheight, imwidth, 1])
masks = sess.run(pred_masks,
feed_dict={train_initial: image_normalized_wn})
self.progress_var.set(j / 2 / num_batches_test * 100)
self.window.update()
# First pass, get the coarse masks, and normalize the image on masks
masks1.append(masks)
# Restore the foreground normalization model from the trained network
saver.restore(sess, './Network/foreground.ckpt')
sess.run(tf.local_variables_initializer())
for j in tqdm(range(0, num_batches_test)):
batch_data = x_test[j]
batch_data_shape = batch_data.shape
image = np.reshape(batch_data,
[batch_data_shape[0], batch_data_shape[1]])
if resize_scale != 1:
image = rescale(image, self.params['scale_ratio'])
# Clip the height and width to be 16-fold
imheight, imwidth = image.shape
imheight = imheight // 16 * 16
imwidth = imwidth // 16 * 16
image = image[:imheight, :imwidth]
# Final pass, foreground normalization to get final masks
image_normalized_fg = foreground_norm(image, masks1[j])
image_normalized_fg = np.reshape(image_normalized_fg,
[1, imheight, imwidth, 1])
# If adding watershed, we save the watershed masks separately
if perform_watershed == 'yes':
masks_watershed = sess.run(
pred_masks_watershed,
feed_dict={train_initial: image_normalized_fg})
if postProcess == 'yes':
masks_watershed = clean_image(masks_watershed)
# Revert the scale to original display
if resize_scale != 1:
masks_watershed = rescale(masks_watershed,
1 / self.params['scale_ratio'])
I8 = (((masks_watershed - masks_watershed.min()) /
(masks_watershed.max() - masks_watershed.min())) *
255).astype(np.uint8)
img = Image.fromarray(I8)
img.save(self.batch_seg_path + x_id[j] + '_masks_watershed.png')
else:
masks = sess.run(pred_masks,
feed_dict={train_initial: image_normalized_fg})
if postProcess == 'yes':
masks = clean_image(masks)
# enable these 2 lines if your want to see the detection result
#image_pil = draw_top_nms_proposals(pred_dict, batch_data, min_score=bbox_min_score, draw_gt=False)
#image_pil.save(str(j)+'_pred.png')
# Revert the scale to original display
if resize_scale != 1:
masks = rescale(masks, 1 / self.params['scale_ratio'])
I8 = (((masks - masks.min()) / (masks.max() - masks.min())) *
255).astype(np.uint8)
img = Image.fromarray(I8)
img.save(self.batch_seg_path + x_id[j] + '_masks.png')
self.progress_var.set(50 + j / 2 / num_batches_test * 100)
self.window.update()
sess.close()
def load_data_train(self, normalization_method='fg'):
"""
Load and normalize the training data from the integrated .pckl file
argument: normalization_method(str): can choose between 'wn'(whole image normalization)
and 'fg'(foreground normalization)
return: the formatted input ready for network
"""
# First load the training image
img_dir = self.train_img_path
imlabel_dir = self.train_label_path
if len(list_files(img_dir, 'png')) > 0:
all_train = list_files(img_dir, 'png')
else:
all_train = list_files(img_dir, 'tif')
if len(list_files(imlabel_dir, 'png')) > 0:
all_train_label = list_files(imlabel_dir, 'png')
else:
all_train_label = list_files(imlabel_dir, 'tif')
if self.usingCL:
print('Computing weight matrix ...')
else:
self.training_results.set('Computing weight matrix ...')
self.window.update()
# Get the data
num_training = len(all_train)
num_training_label = len(all_train_label)
# The number of training images and training labels should be the same
assert num_training == num_training_label
all_train.sort()
all_train_label.sort()
# The training data
x_train = []
# The training label
y_train = []
w_train = []
bbox_train = []
for j in tqdm(range(0, num_training)):
im = Image.open(img_dir + all_train[j])
im = np.asarray(im)
if len(im.shape) > 2:
r, g, b = im[:, :, 0], im[:, :, 1], im[:, :, 2]
im = 0.2989 * r + 0.5870 * g + 0.1140 * b
# fix height and width
height, width = im.shape
width = width // 16 * 16
height = height // 16 * 16
im = im[:height, :width]
iml = Image.open(imlabel_dir + all_train_label[j])
iml = np.asarray(iml)
if len(iml.shape) > 2:
r, g, b = iml[:, :, 0], iml[:, :, 1], iml[:, :, 2]
iml = 0.2989 * r + 0.5870 * g + 0.1140 * b
# fix height and width
height, width = iml.shape
width = width // 16 * 16
height = height // 16 * 16
iml = iml[:height, :width]
# Remove training images with blank labels/annotations
# to avoid dividing by 0
if np.max(iml) > 1:
y_train.append(iml / np.max(iml))
w_train.append(unetwmap(iml / np.max(iml)))
bbox_train.append(bounding_box(iml / np.max(iml)))
x_train.append(im)
elif np.max(iml) == 1:
y_train.append(iml)
w_train.append(unetwmap(iml))
bbox_train.append(bounding_box(iml))
x_train.append(im)
# split the train and validation data 7/8 and 1/8
num_train = len(x_train) // 8 * 7
num_val = len(x_train) - num_train
x_val = x_train[:num_val]
x_train = x_train[num_val:len(x_train)]
y_val = y_train[:num_val]
y_train = y_train[num_val:len(y_train)]
w_val = w_train[:num_val]
w_train = w_train[num_val:len(w_train)]
bbox_val = bbox_train[:num_val]
bbox_train = bbox_train[num_val:len(bbox_train)]
if self.usingCL:
print('Normalizing ...')
else:
self.training_results.set('Normalizing ...')
self.window.update()
if normalization_method == 'wn':
# Normalizing the training data
for i in tqdm(range(len(x_train))):
x_train[i] = whole_image_norm(x_train[i])
# Normalizing the validation data
for i in tqdm(range(len(x_val))):
x_val[i] = whole_image_norm(x_val[i])
if normalization_method == 'fg':
# Normalizing the training data
for i in tqdm(range(len(x_train))):
x_train[i] = foreground_norm(x_train[i], y_train[i])
# Normalizing the validation data: notice it is normalized
# based on whole image norm model predictions
for i in tqdm(range(len(x_val))):
x_val[i] = foreground_norm(x_val[i], self.whole_norm_y_pred[i])
return (x_train, x_val, y_train, y_val, w_train, w_val, bbox_train,
bbox_val)
def test_single_img(params, x_test):
"""input the image, return the segmented mask
Args:
params (dict): the parameters of the network
x_test: the input image in numpy array
"""
# Get the testing parameters
perform_watershed = params['watershed']
bbox_min_score = params['min_score']
nms_thresh = params['nms_threshold']
postProcess = params['postProcess']
# pred_dict and pred_dict_final save all the temp variables
pred_dict_final = {}
train_initial = tf.placeholder(dtype=tf.float32, shape=[1, None, None, 1])
input_shape = tf.shape(train_initial)
input_height = input_shape[1]
input_width = input_shape[2]
im_shape = tf.cast([input_height, input_width], tf.float32)
# number of classes needed to be classified, for our case this equals to 2
# (foreground and background)
nb_classes = 2
# feed the initial image to U-Net, we expect 2 outputs:
# 1. feat_map of shape (?,32,32,1024), which will be passed to the
# region proposal network
# 2. final_logits of shape(?,512,512,2), which is the prediction from U-net
with tf.variable_scope('model_U-Net') as scope:
final_logits, feat_map = UNET(nb_classes, train_initial)
# The final_logits has 2 channels for foreground/background softmax scores,
# then we get prediction with larger score for each pixel
pred_masks = tf.argmax(final_logits, axis=3)
pred_masks = tf.reshape(pred_masks, [input_height, input_width])
pred_masks = tf.to_float(pred_masks)
# Dynamic anchor base size calculated from median cell lengths
base_size = anchor_size(tf.reshape(pred_masks,
[input_height, input_width]))
# scales and ratios are used to generate different anchors
scales = np.array([0.5, 1, 2])
ratios = np.array([0.125, 0.25, 0.5, 1, 2, 4, 8])
# stride is to control how sparse we want to place anchors across the image
# stride = 16 means to place an anchor every 16 pixels on the original image
stride = 16
# Generate the anchor reference with respect to the original image
ref_anchors = generate_anchors_reference(base_size, ratios, scales)
num_ref_anchors = scales.shape[0] * ratios.shape[0]
feat_height = input_height / stride
feat_width = input_width / stride
# Generate all the anchors based on ref_anchors
all_anchors = generate_anchors(ref_anchors, stride,
[feat_height, feat_width])
num_anchors = all_anchors.shape[0]
with tf.variable_scope('model_RPN') as scope:
prediction_dict = RPN(feat_map, num_ref_anchors)
# Get the tensors from the dict
rpn_cls_prob = prediction_dict['rpn_cls_prob']
rpn_bbox_pred = prediction_dict['rpn_bbox_pred']
proposal_prediction = RPNProposal(rpn_cls_prob, rpn_bbox_pred, all_anchors,
im_shape, nms_thresh)
pred_dict_final['all_anchors'] = tf.cast(all_anchors, tf.float32)
prediction_dict['proposals'] = proposal_prediction['proposals']
prediction_dict['scores'] = proposal_prediction['scores']
pred_dict_final['rpn_prediction'] = prediction_dict
scores = pred_dict_final['rpn_prediction']['scores']
proposals = pred_dict_final['rpn_prediction']['proposals']
pred_masks_watershed = tf.to_float(
marker_watershed(scores,
proposals,
pred_masks,
min_score=bbox_min_score))
# start point for testing, and end point for graph
sess = tf.Session()
sess.run(tf.global_variables_initializer())
num_batches_test = len(x_test)
saver = tf.train.Saver()
masks1 = []
# Restore the per-image normalization model from the trained network
saver.restore(sess, './Network/whole_norm.ckpt')
#saver.restore(sess,'./Network/whole_norm_weights_fluorescent/'+str(3)+'.ckpt')
sess.run(tf.local_variables_initializer())
for j in tqdm(range(0, num_batches_test)):
# whole image normalization
batch_data = x_test[j]
batch_data_shape = batch_data.shape
image_normalized_wn = whole_image_norm(batch_data)
image_normalized_wn = np.reshape(
image_normalized_wn,
[1, batch_data_shape[0], batch_data_shape[1], 1])
masks = sess.run(pred_masks,
feed_dict={train_initial: image_normalized_wn})
# First pass, get the coarse masks, and normalize the image on masks
masks1.append(masks)
# Restore the foreground normalization model from the trained network
saver.restore(sess, './Network/foreground.ckpt')
#saver.restore(sess,'./Network/fg_norm_weights_fluorescent/'+str(30)+'.ckpt')
sess.run(tf.local_variables_initializer())
for j in tqdm(range(0, num_batches_test)):
batch_data = x_test[j]
batch_data_shape = batch_data.shape
image = np.reshape(batch_data,
[batch_data_shape[0], batch_data_shape[1]])
# Final pass, foreground normalization to get final masks
image_normalized_fg = foreground_norm(image, masks1[j])
image_normalized_fg = np.reshape(
image_normalized_fg,
[1, batch_data_shape[0], batch_data_shape[1], 1])
# If adding watershed, we save the watershed masks separately
if perform_watershed == 'yes':
masks = sess.run(pred_masks_watershed,
feed_dict={train_initial: image_normalized_fg})
if postProcess == 'yes':
masks = clean_image(masks)
else:
masks = sess.run(pred_masks,
feed_dict={train_initial: image_normalized_fg})
if postProcess == 'yes':
masks = clean_image(masks)
sess.close()
return masks