# Copyright 2022-2024 NXP
diff --git a/tf/backend/utils.py b/tf/backend/utils.py
index cad741e..1160715 100644
--- a/tf/backend/utils.py
+++ b/tf/backend/utils.py
@@ -8,11 +8,15 @@ import torch
 def post_processing(reg_list, cls_list, num_classes, image_size, feature_map_wh_list, min_boxes,
                     center_variance, size_variance,
                     conf_threshold=0.6, nms_max_output_size=100, nms_iou_threshold=0.3, top_k=100):
-    reg_list = [tf.keras.layers.Reshape([-1, 4])(reg) for reg in reg_list]
-    cls_list = [tf.keras.layers.Reshape([-1, num_classes])(cls) for cls in cls_list]
 
-    reg = tf.keras.layers.Concatenate(axis=1)(reg_list)
-    cls = tf.keras.layers.Concatenate(axis=1)(cls_list)
+    shapes = [3600, 600, 160, 60]
+    input_shapes = [(30,40),(15,20),(8,10),(4,5)]
+
+    reg_list = [tf.reshape(reg, shape=(shape,4)) for reg,shape, input_shape in zip(reg_list, shapes, input_shapes)]
+    cls_list = [tf.reshape(cls, shape=(shape, num_classes)) for cls,shape, input_shape in zip(cls_list, shapes, input_shapes)]
+
+    reg = tf.keras.layers.Concatenate(axis=0)(reg_list)
+    cls = tf.keras.layers.Concatenate(axis=0)(cls_list)
 
     # post process
     cls = tf.keras.layers.Softmax(axis=-1)(cls)
@@ -21,23 +25,14 @@ def post_processing(reg_list, cls_list, num_classes, image_size, feature_map_wh_
 
     result = tf.keras.layers.Concatenate(axis=-1)([cls, loc])
 
-    # confidence thresholding
-    mask = conf_threshold < cls[..., 1]
-    result = tf.boolean_mask(tensor=result, mask=mask)
-
     # non-maximum suppression
-    mask = tf.image.non_max_suppression(boxes=result[..., -4:],
-                                        scores=result[..., 1],
+    mask = tf.image.non_max_suppression(boxes=loc,
+                                        scores=cls[:,1],
                                         max_output_size=nms_max_output_size,
                                         iou_threshold=nms_iou_threshold,
                                         name='non_maximum_suppresion')
     result = tf.gather(params=result, indices=mask, axis=0)
 
-    # top-k filtering
-    top_k_value = tf.math.minimum(tf.constant(top_k), tf.shape(result)[0])
-    mask = tf.nn.top_k(result[..., 1], k=top_k_value, sorted=True).indices
-    result = tf.gather(params=result, indices=mask, axis=0)
-
     return result
 
 
diff --git a/tf/convert_tensorflow.py b/tf/convert_tensorflow.py
index ca6b99b..6f4f185 100644
--- a/tf/convert_tensorflow.py
+++ b/tf/convert_tensorflow.py
@@ -1,9 +1,9 @@
 import argparse
 import sys
 
-from tf.backend.utils import load_weight
-from tf.model.rfb_320 import create_rfb_net
-from tf.model.slim_320 import create_slim_net
+from backend.utils import load_weight
+from model.rfb_320 import create_rfb_net
+from model.slim_320 import create_slim_net
 
 parser = argparse.ArgumentParser(
     description='convert model')
diff --git a/tf/model/rfb_320.py b/tf/model/rfb_320.py
index 387edd0..9ea64f5 100644
--- a/tf/model/rfb_320.py
+++ b/tf/model/rfb_320.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
 
-from tf.backend.op import conv_bn, conv_dw, basic_rfb, separable_conv
-from tf.backend.utils import post_processing
+from backend.op import conv_bn, conv_dw, basic_rfb, separable_conv
+from backend.utils import post_processing
 
 conf_threshold = 0.6
 nms_iou_threshold = 0.3
diff --git a/tf/model/slim_320.py b/tf/model/slim_320.py
index 5b1b27a..519d0a0 100644
--- a/tf/model/slim_320.py
+++ b/tf/model/slim_320.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
 
-from tf.backend.op import conv_bn, conv_dw, separable_conv
-from tf.backend.utils import post_processing
+from backend.op import conv_bn, conv_dw, separable_conv
+from backend.utils import post_processing
 
 conf_threshold = 0.6
 nms_iou_threshold = 0.3