OpenVINO NonMaxSuppression -> TensorFlow NonMaxSuppressionV5

import os
import warnings
os.environ['TF_CPP_MIN_LOG_LEVEL']='3'
warnings.simplefilter(action='ignore', category=FutureWarning)
warnings.simplefilter(action='ignore', category=Warning)
import tensorflow as tf
import numpy as np
np.set_printoptions(edgeitems=90)
np.random.seed(seed=32)

batch = 1
classes = 20 #20
boxes = 5 #19206
max_output_size = 3 #Top 3 conf boexes
iou_threshold = 0.5
score_threshold = 0.0
soft_nms_sigma = 0.0

a = np.random.rand(batch, classes, boxes) # batch, classes, boxes
print(a)
print(a.shape)

b = a.transpose(0,2,1) # [batch, boxes, classes]
c = tf.constant(a.transpose(0,2,1), dtype=tf.float32)

# Find top elements
top, top_idx = tf.nn.top_k(input=c, k=1, sorted=False)
print('@@@@@@@@@@@@@@@@@@@@@@@ top - boxes - conf')
top_flat = tf.reshape(top, [-1])
print(top_flat)
print('@@@@@@@@@@@@@@@@@@@@@@@ top_idx - boxes - class_id')
top_idx = tf.reshape(top_idx, [-1])
print(top_idx)

output_size = tf.math.minimum(boxes, max_output_size)
print('@@@@@@@@@@@@@@@@@@@@@@@ output_size')
print(output_size)

#==============================================================
# openvino   = [y1, x1, y2, x2]
# tensorflow = [y1, x1, y2, x2]
xyxy = 4
d = np.random.rand(boxes, xyxy)
print('@@@@@@@@@@@@@@@@@@@@@@@ boxes - xy')
print(d)
print(d.shape)

aaa = tf.constant(d, dtype=tf.float32)

#==============================================================
selected_indices, selected_scores, valid_outputs = \
    tf.raw_ops.NonMaxSuppressionV5(boxes=aaa,
                                    scores=top_flat,
                                    max_output_size=output_size,
                                    iou_threshold=iou_threshold,
                                    score_threshold=score_threshold,
                                    soft_nms_sigma=soft_nms_sigma,
                                    pad_to_max_output_size=True)
print('@@@@@@@@@@@@@@@@@@@@@@@ NonMaxSuppressionV5')
print('selected_indices')
print(selected_indices)
print('selected_scores')
print(selected_scores)

selected_boxes = tf.gather(d, selected_indices)
selected_class_idx = tf.gather(top_idx, selected_indices)

#============================================================
print('selected_boxes')
print(selected_boxes)
print('selected_scores')
print(selected_scores)
print('selected_class_idx')
print(selected_class_idx)
print('boxes_count')
print(output_size)