example-apps/PDD/pcb-defect-detection/libs/box_utils/show_box_in_tensor.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import absolute_import
   4 from __future__ import division
   5 from __future__ import print_function
   6
   7 import tensorflow as tf
   8 import numpy as np
   9 import cv2
  10 from libs.label_name_dict.label_dict import LABEl_NAME_MAP
  11
  12 from libs.configs import cfgs
  13
  14 from libs.box_utils import draw_box_in_img
  15
  16 def only_draw_boxes(img_batch, boxes):
  17
  18     boxes = tf.stop_gradient(boxes)
  19     img_tensor = tf.squeeze(img_batch, 0)
  20     img_tensor = tf.cast(img_tensor, tf.float32)
  21     labels = tf.ones(shape=(tf.shape(boxes)[0], ), dtype=tf.int32) * draw_box_in_img.ONLY_DRAW_BOXES
  22     scores = tf.zeros_like(labels, dtype=tf.float32)
  23     img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores,
  24                                        inp=[img_tensor, boxes, labels, scores],
  25                                        Tout=tf.uint8)
  26     img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))  # [batch_size, h, w, c]
  27
  28     return img_tensor_with_boxes
  29
  30 def draw_boxes_with_scores(img_batch, boxes, scores):
  31
  32     boxes = tf.stop_gradient(boxes)
  33     scores = tf.stop_gradient(scores)
  34
  35     img_tensor = tf.squeeze(img_batch, 0)
  36     img_tensor = tf.cast(img_tensor, tf.float32)
  37     labels = tf.ones(shape=(tf.shape(boxes)[0],), dtype=tf.int32) * draw_box_in_img.ONLY_DRAW_BOXES_WITH_SCORES
  38     img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores,
  39                                        inp=[img_tensor, boxes, labels, scores],
  40                                        Tout=[tf.uint8])
  41     img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))
  42     return img_tensor_with_boxes
  43
  44 def draw_boxes_with_categories(img_batch, boxes, labels):
  45     boxes = tf.stop_gradient(boxes)
  46
  47     img_tensor = tf.squeeze(img_batch, 0)
  48     img_tensor = tf.cast(img_tensor, tf.float32)
  49     scores = tf.ones(shape=(tf.shape(boxes)[0],), dtype=tf.float32)
  50     img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores,
  51                                        inp=[img_tensor, boxes, labels, scores],
  52                                        Tout=[tf.uint8])
  53     img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))
  54     return img_tensor_with_boxes
  55
  56 def draw_boxes_with_categories_and_scores(img_batch, boxes, labels, scores):
  57     boxes = tf.stop_gradient(boxes)
  58     scores = tf.stop_gradient(scores)
  59
  60     img_tensor = tf.squeeze(img_batch, 0)
  61     img_tensor = tf.cast(img_tensor, tf.float32)
  62     img_tensor_with_boxes = tf.py_func(draw_box_in_img.draw_boxes_with_label_and_scores,
  63                                        inp=[img_tensor, boxes, labels, scores],
  64                                        Tout=[tf.uint8])
  65     img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))
  66     return img_tensor_with_boxes
  67
  68 if __name__ == "__main__":
  69     print (1)
  70