example-apps/PDD/pcb-defect-detection/libs/losses/losses.py

   1 # -*- coding: utf-8 -*-
   2 """
   3 @author: jemmy li
   4 @contact: zengarden2009@gmail.com
   5 """
   6 from __future__ import absolute_import
   7 from __future__ import division
   8 from __future__ import print_function
   9
  10 import tensorflow as tf
  11
  12
  13 def _smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=1.0):
  14     '''
  15
  16     :param bbox_pred: [-1, 4] in RPN. [-1, cls_num+1, 4] in Fast-rcnn
  17     :param bbox_targets: shape is same as bbox_pred
  18     :param sigma:
  19     :return:
  20     '''
  21     sigma_2 = sigma**2
  22
  23     box_diff = bbox_pred - bbox_targets
  24
  25     abs_box_diff = tf.abs(box_diff)
  26
  27     smoothL1_sign = tf.stop_gradient(
  28         tf.to_float(tf.less(abs_box_diff, 1. / sigma_2)))
  29     loss_box = tf.pow(box_diff, 2) * (sigma_2 / 2.0) * smoothL1_sign \
  30                + (abs_box_diff - (0.5 / sigma_2)) * (1.0 - smoothL1_sign)
  31     return loss_box
  32
  33 def smooth_l1_loss_rpn(bbox_pred, bbox_targets, label, sigma=1.0):
  34     '''
  35
  36     :param bbox_pred: [-1, 4]
  37     :param bbox_targets: [-1, 4]
  38     :param label: [-1]
  39     :param sigma:
  40     :return:
  41     '''
  42     value = _smooth_l1_loss_base(bbox_pred, bbox_targets, sigma=sigma)
  43     value = tf.reduce_sum(value, axis=1)  # to sum in axis 1
  44     rpn_positive = tf.where(tf.greater(label, 0))
  45
  46     # rpn_select = tf.stop_gradient(rpn_select) # to avoid
  47     selected_value = tf.gather(value, rpn_positive)
  48     non_ignored_mask = tf.stop_gradient(
  49         1.0 - tf.to_float(tf.equal(label, -1)))  # positve is 1.0 others is 0.0
  50
  51     bbox_loss = tf.reduce_sum(selected_value) / tf.maximum(1.0, tf.reduce_sum(non_ignored_mask))
  52
  53     return bbox_loss
  54
  55 def smooth_l1_loss_rcnn(bbox_pred, bbox_targets, label, num_classes, sigma=1.0):
  56     '''
  57
  58     :param bbox_pred: [-1, (cfgs.CLS_NUM +1) * 4]
  59     :param bbox_targets:[-1, (cfgs.CLS_NUM +1) * 4]
  60     :param label:[-1]
  61     :param num_classes:
  62     :param sigma:
  63     :return:
  64     '''
  65
  66     outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))
  67
  68     bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
  69     bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])
  70
  71     value = _smooth_l1_loss_base(bbox_pred,
  72                                  bbox_targets,
  73                                  sigma=sigma)
  74     value = tf.reduce_sum(value, 2)
  75     value = tf.reshape(value, [-1, num_classes])
  76
  77     inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
  78                              depth=num_classes, axis=1)
  79
  80     inside_mask = tf.stop_gradient(
  81         tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))
  82
  83     normalizer = tf.to_float(tf.shape(bbox_pred)[0])
  84     bbox_loss = tf.reduce_sum(
  85         tf.reduce_sum(value * inside_mask, 1)*outside_mask) / normalizer
  86
  87     return bbox_loss
  88
  89 def sum_ohem_loss(cls_score, label, bbox_pred, bbox_targets,
  90                   num_classes, num_ohem_samples=256, sigma=1.0):
  91
  92     '''
  93     :param cls_score: [-1, cls_num+1]
  94     :param label: [-1]
  95     :param bbox_pred: [-1, 4*(cls_num+1)]
  96     :param bbox_targets: [-1, 4*(cls_num+1)]
  97     :param num_ohem_samples: 256 by default
  98     :param num_classes: cls_num+1
  99     :param sigma:
 100     :return:
 101     '''
 102     cls_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=cls_score, labels=label)  # [-1, ]
 103     # cls_loss = tf.Print(cls_loss, [tf.shape(cls_loss)], summarize=10, message='CLS losss shape ****')
 104
 105     outside_mask = tf.stop_gradient(tf.to_float(tf.greater(label, 0)))
 106     bbox_pred = tf.reshape(bbox_pred, [-1, num_classes, 4])
 107     bbox_targets = tf.reshape(bbox_targets, [-1, num_classes, 4])
 108
 109     value = _smooth_l1_loss_base(bbox_pred,
 110                                  bbox_targets,
 111                                  sigma=sigma)
 112     value = tf.reduce_sum(value, 2)
 113     value = tf.reshape(value, [-1, num_classes])
 114
 115     inside_mask = tf.one_hot(tf.reshape(label, [-1, 1]),
 116                              depth=num_classes, axis=1)
 117
 118     inside_mask = tf.stop_gradient(
 119         tf.to_float(tf.reshape(inside_mask, [-1, num_classes])))
 120     loc_loss = tf.reduce_sum(value * inside_mask, 1)*outside_mask
 121     # loc_loss = tf.Print(loc_loss, [tf.shape(loc_loss)], summarize=10, message='loc_loss shape***')
 122
 123     sum_loss = cls_loss + loc_loss
 124
 125     num_ohem_samples = tf.stop_gradient(tf.minimum(num_ohem_samples, tf.shape(sum_loss)[0]))
 126     _, top_k_indices = tf.nn.top_k(sum_loss, k=num_ohem_samples)
 127
 128     cls_loss_ohem = tf.gather(cls_loss, top_k_indices)
 129     cls_loss_ohem = tf.reduce_mean(cls_loss_ohem)
 130
 131     loc_loss_ohem = tf.gather(loc_loss, top_k_indices)
 132     normalizer = tf.to_float(num_ohem_samples)
 133     loc_loss_ohem = tf.reduce_sum(loc_loss_ohem) / normalizer
 134
 135     return cls_loss_ohem, loc_loss_ohem