def resnet_fpn_model_fn(features, labels, mode, params):
"""parameter definition part"""
is_train = (mode == tf.estimator.ModeKeys.TRAIN)
resnet_num_blocks = params["RESNET_NUM_BLOCKS"]
num_anchors = params["num_anchors"]
head_dim = params["head_dim"]
resolution = params["resolution"]
num_classes = params["num_classes"]
bbox_reg_weights = params["bbox_reg_weights"]
weight_decay = params["weight_decay"]
learning_rate = params["learning_rate"]
lr_schedule = params["lr_schedule"]
"""model definition part"""
image = image_preprocess(features["image"])
c2345 = resnet_fpn_backbone(image, resnet_num_blocks, is_train)
p23456 = fpn_model(c2345)
image_shape2d = tf.shape(image)[1:3]
all_anchors_fpn = tf_get_all_anchors_fpn()
model_rpn_head = RPNHead(_C.FPN.NUM_CHANNEL, len(_C.RPN.ANCHOR_RATIOS))
rpn_outputs = [model_rpn_head(pi) for pi in p23456]
multilevel_label_logits = [k[0] for k in rpn_outputs]
multilevel_box_logits = [k[1] for k in rpn_outputs]
# debug_op = tf.print({"debug_inf": tf.convert_to_tensor("now in here")})
# with tf.control_dependencies([debug_op]):
# image_shape2d = tf.identity(image_shape2d)
if mode != tf.estimator.ModeKeys.PREDICT:
multilevel_anchors = [
RPNAnchors(all_anchors_fpn[i],
features['anchor_labels_lvl{}'.format(i + 2)],
features['anchor_boxes_lvl{}'.format(i + 2)])
for i in range(len(all_anchors_fpn))
]
else:
multilevel_anchors = [
RPNAnchors(all_anchors_fpn[i], None, None)
for i in range(len(all_anchors_fpn))
]
slice_feature_and_anchors(p23456, multilevel_anchors)
# Multi-Level RPN Proposals
multilevel_pred_boxes = [
anchor.decode_logits(logits)
for anchor, logits in zip(multilevel_anchors, multilevel_box_logits)
]
proposal_boxes, proposal_scores = generate_fpn_proposals(
multilevel_pred_boxes, multilevel_label_logits, image_shape2d,
is_train)
proposals = BoxProposals(proposal_boxes)
gt_boxes = None
gt_labels = None
if mode != tf.estimator.ModeKeys.PREDICT:
losses = multilevel_rpn_losses(multilevel_anchors,
multilevel_label_logits,
multilevel_box_logits)
gt_boxes = features['boxes']
gt_labels = features['gt_labels']
proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes,
gt_labels)
fastrcnn_head_func = getattr(model_frcnn, _C.FPN.FRCNN_HEAD_FUNC)
if not _C.FPN.CASCADE:
roi_feature_fastrcnn = multilevel_roi_align(p23456[:4],
proposals.boxes, 7)
head_feature = fastrcnn_head_func(roi_feature_fastrcnn)
fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(
head_feature, num_classes)
fastrcnn_head = FastRCNNHead(
proposals, fastrcnn_box_logits, fastrcnn_label_logits, gt_boxes,
tf.convert_to_tensor(bbox_reg_weights, tf.float32))
else:
def roi_func(boxes):
return multilevel_roi_align(p23456[:4], boxes, 7)
fastrcnn_head = CascadeRCNNHead(proposals, roi_func,
fastrcnn_head_func,
(gt_boxes, gt_labels), image_shape2d,
num_classes,
mode != tf.estimator.ModeKeys.PREDICT)
decoded_boxes = fastrcnn_head.decoded_output_boxes()
decoded_boxes = clip_boxes(decoded_boxes, image_shape2d)
label_scores = fastrcnn_head.output_scores()
# final_boxes, final_scores, final_labels = fastrcnn_predictions(decoded_boxes, label_scores)
final_boxes, final_scores, final_labels, valid_detections = fastrcnn_predictions_v2(
decoded_boxes, label_scores)
global_step = tf.train.get_or_create_global_step()
if mode != tf.estimator.ModeKeys.PREDICT:
all_losses = fastrcnn_head.losses()
trainable_weights = tf.trainable_variables()
weight_loss = 0.0
for i, ele in enumerate(trainable_weights):
if re.search('.*/kernel', ele.name):
weight_loss += tf.reduce_sum(tf.square(ele) * weight_decay)
# print_op = tf.print({'rpn_loss': tf.add_n(losses),
# 'frcnn_loss': tf.add_n(all_losses)})
# with tf.control_dependencies([print_op]):
total_cost = tf.add_n(losses + all_losses, "total_cost")
total_cost = tf.add(total_cost, weight_loss, 'all_total_cost')
if is_train:
learning_rate = tf.train.piecewise_constant(
global_step,
lr_schedule,
values=[tf.convert_to_tensor(0.01 * 0.33 * 0.375, tf.float32)]
+ [learning_rate * (0.1**i) for i in range(len(lr_schedule))])
tf.summary.scalar("learning_rate", learning_rate)
opt = tf.train.MomentumOptimizer(learning_rate, 0.9)
# opt = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
print(update_ops)
with tf.control_dependencies(update_ops):
train_op = opt.minimize(total_cost, global_step)
return tf.estimator.EstimatorSpec(mode,
loss=total_cost,
train_op=train_op)
else:
return tf.estimator.EstimatorSpec(mode, loss=total_cost)
else:
predictions = {
'boxes': final_boxes[0, :valid_detections[0]],
'labels': final_labels[0, :valid_detections[0]],
'scores': final_scores[0, :valid_detections[0]],
# 'image': features['image'],
# 'original_image': features['original_image'],
'h_pre': features['h_pre'],
'w_pre': features['w_pre'],
'h_now': features['h_now'],
'w_now': features['w_now'],
'scale': features['scale'],
'image_id': features.get('image_id', tf.convert_to_tensor(0)),
'valid_detection': valid_detections
}
# predictions = {'boxes': final_boxes,
# 'labels': final_labels,
# 'scores': final_scores,
# 'image': features['image'],
# 'original_image': features['original_image'],
# 'h_pre': features['h_pre'],
# 'w_pre': features['w_pre'],
# 'h_now': features['h_now'],
# 'w_now': features['w_now'],
# 'scale': features['scale'],
# 'image_id': features.get('image_id', tf.convert_to_tensor(0))
# }
return tf.estimator.EstimatorSpec(mode, predictions)
def resnet_c4_model_fn(features, labels, mode, params):
"""parameter defination part"""
is_train = (mode == tf.estimator.ModeKeys.TRAIN)
resnet_num_blocks = params["RESNET_NUM_BLOCKS"]
num_anchors = params["num_anchors"]
head_dim = params["head_dim"]
resolution = params["resolution"]
num_classes = params["num_classes"]
bbox_reg_weights = params["bbox_reg_weights"]
weight_decay = params["weight_decay"]
learning_rate = params["learning_rate"]
lr_schedule = params["lr_schedule"]
"""model definition part"""
image = image_preprocess(features['image'])
featuremap = resnet_c4_backbone(image, resnet_num_blocks[:3], is_train)
image_shape2d = tf.shape(image)[1:3]
rpn_label_logits, rpn_box_logits = rpn_head(featuremap, head_dim,
num_anchors)
if mode != tf.estimator.ModeKeys.PREDICT:
anchors = RPNAnchors(tf_get_all_anchors(), features['anchor_labels'],
features['anchor_boxes'])
else:
anchors = RPNAnchors(tf_get_all_anchors(), None, None)
anchors = anchors.narrow_to(featuremap)
pred_boxes_decoded = anchors.decode_logits(rpn_box_logits) # x1y1x2y2
proposals, proposal_scores = generate_rpn_proposals(
tf.reshape(pred_boxes_decoded, [-1, 4]),
tf.reshape(rpn_label_logits,
[-1]), image_shape2d, _C.RPN.TRAIN_PRE_NMS_TOPK
if mode == tf.estimator.ModeKeys.TRAIN else _C.RPN.TEST_PRE_NMS_TOPK,
_C.RPN.TRAIN_POST_NMS_TOPK
if mode == tf.estimator.ModeKeys.TRAIN else _C.RPN.TEST_POST_NMS_TOPK)
# rpn_size = tf.shape(proposals)[0]
# rpn_boxes = tf.gather(proposals, tf.where(tf.greater(proposals, 0.5)))
proposals = BoxProposals(proposals)
if mode != tf.estimator.ModeKeys.PREDICT:
rpn_loss = rpn_losses(anchors.gt_labels, anchors.encoded_gt_boxes(),
rpn_label_logits, rpn_box_logits)
# targets = [features[k] for k in ['boxes', 'gt_labels', 'gt_masks'] if k in features.keys()]
# gt_boxes, gt_labels, *_ = targets
gt_boxes = features['boxes']
gt_labels = features['gt_labels']
proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes,
gt_labels)
boxes_on_featuremap = proposals.boxes * (1.0 / _C.RPN.ANCHOR_STRIDE)
roi_resized = roi_align(featuremap, boxes_on_featuremap, resolution)
feature_fastrcnn = resnet_conv5(roi_resized, resnet_num_blocks[-1],
is_train)
# Keep C5 feature to be shared with mask branch
feature_gap = keras.layers.GlobalAveragePooling2D()(feature_fastrcnn)
fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(
feature_gap, num_classes)
bbox_reg_weights_tensor = tf.convert_to_tensor(bbox_reg_weights,
tf.float32)
if mode != tf.estimator.ModeKeys.PREDICT:
fastrcnn_head = FastRCNNHead(proposals, fastrcnn_box_logits,
fastrcnn_label_logits, gt_boxes,
bbox_reg_weights_tensor)
all_loss = fastrcnn_head.losses()
label_scores = tf.nn.softmax(fastrcnn_label_logits)
decoded_boxes = decoded_output_boxes(proposals, num_classes,
fastrcnn_box_logits,
bbox_reg_weights_tensor)
decoded_boxes = clip_boxes(decoded_boxes, image_shape2d)
final_boxes, final_scores, final_labels, valid_detections = fastrcnn_predictions_v2(
decoded_boxes, label_scores)
# final_boxes, final_scores, final_labels = fastrcnn_predictions(decoded_boxes, label_scores)
global_step = tf.train.get_or_create_global_step()
if mode != tf.estimator.ModeKeys.PREDICT:
trainable_weights = tf.trainable_variables()
weight_loss = 0.0
for i, ele in enumerate(trainable_weights):
if re.search('.*/kernel', ele.name):
weight_loss += tf.reduce_sum(tf.square(ele) * weight_decay)
total_cost = tf.add_n(rpn_loss + all_loss, 'total_cost')
tf.summary.scalar('total_cost', total_cost)
if is_train:
learning_rate = tf.train.piecewise_constant(
global_step,
lr_schedule,
values=[tf.convert_to_tensor(0.01 * 0.33, tf.float32)] +
[learning_rate * (0.1**i) for i in range(len(lr_schedule))])
opt = tf.train.MomentumOptimizer(learning_rate, 0.9)
# opt = tf.train.AdamOptimizer(learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops):
train_op = opt.minimize(total_cost, global_step)
return tf.estimator.EstimatorSpec(mode,
loss=total_cost,
train_op=train_op)
else:
return tf.estimator.EstimatorSpec(mode, loss=total_cost)
else:
predictions = {
'boxes': final_boxes[0, :valid_detections[0]],
'labels': final_labels[0, :valid_detections[0]],
'scores': final_scores[0, :valid_detections[0]],
'image': features['image'],
# 'rpn_boxes': rpn_boxes,
# 'rpn_size': rpn_size,
'valid_detection': valid_detections
}
return tf.estimator.EstimatorSpec(mode, predictions)
tf.reshape(pred_boxes_decoded, [-1, 4]),
tf.reshape(rpn_label_logits, [-1]),
image_shape2d,
_C.RPN.TRAIN_PRE_NMS_TOPK if is_train else _C.RPN.TEST_PRE_NMS_TOPK,
_C.RPN.TRAIN_POST_NMS_TOPK if is_train else _C.RPN.TEST_POST_NMS_TOPK)
proposals = BoxProposals(proposals)
if is_train:
rpn_loss = rpn_losses(anchors.gt_labels, anchors.encoded_gt_boxes(), rpn_label_logits, rpn_box_logits)
# targets = [features[k] for k in ['boxes', 'gt_labels', 'gt_masks'] if k in features.keys()]
# gt_boxes, gt_labels, *_ = targets
gt_boxes = features['boxes']
gt_labels = features['gt_labels']
proposals = sample_fast_rcnn_targets(proposals.boxes, gt_boxes, gt_labels)
boxes_on_featuremap = proposals.boxes * (1.0 / _C.RPN.ANCHOR_STRIDE)
roi_resized = roi_align(featuremap, boxes_on_featuremap, resolution)
feature_fastrcnn = resnet_conv5(roi_resized, resnet_num_blocks[-1], is_train)
# Keep C5 feature to be shared with mask branch
feature_gap = keras.layers.GlobalAveragePooling2D()(feature_fastrcnn)
fastrcnn_label_logits, fastrcnn_box_logits = fastrcnn_outputs(feature_gap, num_classes)
bbox_reg_weights_tensor = tf.convert_to_tensor(bbox_reg_weights, tf.float32)
if is_train:
fastrcnn_head = FastRCNNHead(proposals, fastrcnn_box_logits, fastrcnn_label_logits,
gt_boxes, bbox_reg_weights_tensor)
all_loss = fastrcnn_head.losses()
label_scores = tf.nn.softmax(fastrcnn_label_logits)
decoded_boxes = decoded_output_boxes(proposals, num_classes, fastrcnn_box_logits, bbox_reg_weights_tensor)
decoded_boxes = clip_boxes(decoded_boxes, image_shape2d)
final_boxes, final_scores, final_labels = fastrcnn_predictions(decoded_boxes, label_scores)
print("final_boxes", final_boxes)
print("final_scores", final_scores)
print("final_labels", final_labels)