Implemented resize before debayering.

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"],
     )