1 # -*- coding: utf-8 -*-
3 from __future__ import absolute_import
4 from __future__ import division
5 from __future__ import print_function
7 import tensorflow as tf
10 from libs.label_name_dict.label_dict import LABEl_NAME_MAP
12 from libs.configs import cfgs
14 from libs.box_utils import draw_box_in_img
16 def only_draw_boxes(img_batch, boxes):
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],
26 img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch)) # [batch_size, h, w, c]
28 return img_tensor_with_boxes
30 def draw_boxes_with_scores(img_batch, boxes, scores):
32 boxes = tf.stop_gradient(boxes)
33 scores = tf.stop_gradient(scores)
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],
41 img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))
42 return img_tensor_with_boxes
44 def draw_boxes_with_categories(img_batch, boxes, labels):
45 boxes = tf.stop_gradient(boxes)
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],
53 img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))
54 return img_tensor_with_boxes
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)
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],
65 img_tensor_with_boxes = tf.reshape(img_tensor_with_boxes, tf.shape(img_batch))
66 return img_tensor_with_boxes
68 if __name__ == "__main__":