diff --git a/extras/mmdeploy/make_extra_graphs_tf.py b/extras/mmdeploy/make_extra_graphs_tf.py index e5eb135..3cf1d27 100644 --- a/extras/mmdeploy/make_extra_graphs_tf.py +++ b/extras/mmdeploy/make_extra_graphs_tf.py @@ -90,7 +90,6 @@ class BayerToRGB(tf.keras.layers.Layer): ) def call(self, img): - img = tf.cast(img, tf.float32) H, W = tf.shape(img)[1], tf.shape(img)[2] # Pad the image @@ -129,11 +128,36 @@ class BayerToRGB(tf.keras.layers.Layer): # ================================================================================================== +def bayer_resize(img, size): + """Resize a Bayer image by splitting color channels""" + + # Split the image into 4 channels + r = img[:, 0::2, 0::2, 0] + g1 = img[:, 0::2, 1::2, 0] + g2 = img[:, 1::2, 0::2, 0] + b = img[:, 1::2, 1::2, 0] + bsplit = tf.stack([r, g1, g2, b], axis=-1) + + # Resize the image + # Make sure the target size is divisible by 2 + size = (size[0] // 2, size[1] // 2) + bsized = tf.image.resize(bsplit, size=size, method="bilinear") + + # Create a bayer image again + img = tf.nn.depth_to_space(bsized, block_size=2) + + return img + + +# ================================================================================================== + + class Letterbox(tf.keras.layers.Layer): def __init__(self, target_size, fill_value=128): """Resize and pad image while keeping aspect ratio""" super(Letterbox, self).__init__() + self.b2rgb = BayerToRGB() self.target_size = target_size self.fill_value = fill_value @@ -146,6 +170,8 @@ class Letterbox(tf.keras.layers.Layer): new_h = tf.round(tf.cast(img_h, scale.dtype) * scale) new_w = tf.cast(new_w, tf.int32) new_h = tf.cast(new_h, tf.int32) + new_w = new_w - (new_w % 2) + new_h = new_h - (new_h % 2) pad_w = target_w - new_w pad_h = target_h - new_h @@ -160,9 +186,9 @@ class Letterbox(tf.keras.layers.Layer): def call(self, img): paddings, _, (nw, nh) = self.calc_params(tf.shape(img)) - # Resize the image - img = tf.cast(img, tf.float32) - img = tf.image.resize(img, size=(nh, nw), method="bilinear") + # Resize the image and convert to RGB + img = bayer_resize(img, (nh, nw)) + img = self.b2rgb(img) # Pad the image pad_top, pad_bottom, pad_left, pad_right = paddings @@ -182,13 +208,14 @@ class Letterbox(tf.keras.layers.Layer): class DetPreprocess(tf.keras.layers.Layer): def __init__(self, target_size, fill_value=114): super(DetPreprocess, self).__init__() - - self.b2rgb = BayerToRGB() self.letterbox = Letterbox(target_size, fill_value) def call(self, img): - # img: tf.Tensor of shape [batch, H, W, C], dtype=tf.uint8 - img = self.b2rgb(img) + """img: tf.Tensor of shape [batch, H, W, C], dtype=tf.uint8""" + + # Cast to float32 since TensorRT does not support uint8 layers + img = tf.cast(img, tf.float32) + img = self.letterbox(img) return img @@ -239,6 +266,7 @@ def main(): input_signature, opset=11, output_path=base_path + "det_preprocess.onnx", + target=["tensorrt"], )