-
Notifications
You must be signed in to change notification settings - Fork 6
/
nms_tf_test.py
102 lines (88 loc) · 2.48 KB
/
nms_tf_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
import tensorflow as tf
import numpy as np
from pprint import pprint
np.random.seed(0)
# dummy_input = np.ones([1,10], dtype=np.float32)
BOXES=18900
# Create a model
boxes = tf.keras.layers.Input(
shape=[
BOXES,
1,
4,
],
batch_size=1,
dtype=tf.float32,
)
scores = tf.keras.layers.Input(
shape=[
BOXES,
1,
],
batch_size=1,
dtype=tf.float32,
)
boxes_non_batch = tf.squeeze(boxes)
x1 = boxes_non_batch[:,0][:,np.newaxis]
y1 = boxes_non_batch[:,1][:,np.newaxis]
x2 = boxes_non_batch[:,2][:,np.newaxis]
y2 = boxes_non_batch[:,3][:,np.newaxis]
boxes_y1x1y2x2 = tf.concat([y1,x1,y2,x2], axis=1)
scores_non_batch = tf.squeeze(scores)
selected_indices = tf.image.non_max_suppression(
boxes=boxes_y1x1y2x2,
scores=scores_non_batch,
max_output_size=100,
iou_threshold=0.5,
score_threshold=float('-inf'),
)
# selected_indices = tf.raw_ops.NonMaxSuppressionV3(
# boxes=boxes[0],
# scores=scores[0],
# max_output_size=100,
# iou_threshold=0.5,
# score_threshold=float('-inf'),
# )
# def NonMaxSuppressionV3_(boxes, scores, max_output_size: int, iou_threshold, score_threshold):
# selected_indices = \
# tf.raw_ops.NonMaxSuppressionV3(
# boxes=boxes,
# scores=scores,
# max_output_size=max_output_size,
# iou_threshold=iou_threshold,
# score_threshold=score_threshold
# )
# return selected_indices
# selected_indices = \
# tf.keras.layers.Lambda(
# NonMaxSuppressionV3_,
# arguments={
# 'scores': scores[0],
# 'max_output_size': 100,
# 'iou_threshold': 0.5,
# 'score_threshold': float('-inf'),
# }
# )(boxes[0])
selected_boxes = tf.gather(
boxes_non_batch,
selected_indices
)
selected_scores = tf.gather(
scores_non_batch,
selected_indices
)
selected_scores = tf.expand_dims(selected_scores, axis=1)
outputs = tf.concat([selected_boxes, selected_scores], axis=1)
model = tf.keras.models.Model(inputs=[boxes,scores], outputs=[outputs])
model.summary()
output_path = 'saved_model_postprocess'
tf.saved_model.save(model, output_path)
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.target_spec.supported_ops = [
tf.lite.OpsSet.TFLITE_BUILTINS,
tf.lite.OpsSet.SELECT_TF_OPS
]
tflite_model = converter.convert()
open(f"{output_path}/test.tflite", "wb").write(tflite_model)