From 047c2408d20c2ed7765d7a0216fe064eb0ca166a Mon Sep 17 00:00:00 2001
From: Blake Blackshear <blakeb@blakeshome.com>
Date: Sat, 26 Nov 2022 19:15:47 -0600
Subject: [PATCH] fix support for yuv models (#4523)

---
 frigate/util.py  | 52 ++++++++++++++++++++++++++++++++++++++++++++++++
 frigate/video.py |  4 ++--
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/frigate/util.py b/frigate/util.py
index 05e5fec29..e6a5df4e6 100755
--- a/frigate/util.py
+++ b/frigate/util.py
@@ -379,6 +379,47 @@ def yuv_crop_and_resize(frame, region, height=None):
     return yuv_cropped_frame
 
 
+def yuv_to_3_channel_yuv(yuv_frame):
+    height = yuv_frame.shape[0] // 3 * 2
+    width = yuv_frame.shape[1]
+
+    # flatten the image into array
+    yuv_data = yuv_frame.ravel()
+
+    # create a numpy array to hold all the 3 chanel yuv data
+    all_yuv_data = np.empty((height, width, 3), dtype=np.uint8)
+
+    y_count = height * width
+    uv_count = y_count // 4
+
+    # copy the y_channel
+    all_yuv_data[:, :, 0] = yuv_data[0:y_count].reshape((height, width))
+    # copy the u channel doubling each dimension
+    all_yuv_data[:, :, 1] = np.repeat(
+        np.reshape(
+            np.repeat(yuv_data[y_count : y_count + uv_count], repeats=2, axis=0),
+            (height // 2, width),
+        ),
+        repeats=2,
+        axis=0,
+    )
+    # copy the v channel doubling each dimension
+    all_yuv_data[:, :, 2] = np.repeat(
+        np.reshape(
+            np.repeat(
+                yuv_data[y_count + uv_count : y_count + uv_count + uv_count],
+                repeats=2,
+                axis=0,
+            ),
+            (height // 2, width),
+        ),
+        repeats=2,
+        axis=0,
+    )
+
+    return all_yuv_data
+
+
 def copy_yuv_to_position(
     destination_frame,
     destination_offset,
@@ -497,6 +538,17 @@ def copy_yuv_to_position(
         )
 
 
+def yuv_region_2_yuv(frame, region):
+    try:
+        # TODO: does this copy the numpy array?
+        yuv_cropped_frame = yuv_crop_and_resize(frame, region)
+        return yuv_to_3_channel_yuv(yuv_cropped_frame)
+    except:
+        print(f"frame.shape: {frame.shape}")
+        print(f"region: {region}")
+        raise
+
+
 def yuv_region_2_rgb(frame, region):
     try:
         # TODO: does this copy the numpy array?
diff --git a/frigate/video.py b/frigate/video.py
index 06daa8ea2..bbb5dd866 100755
--- a/frigate/video.py
+++ b/frigate/video.py
@@ -29,9 +29,9 @@ from frigate.util import (
     intersection,
     intersection_over_union,
     listen,
-    yuv_crop_and_resize,
     yuv_region_2_rgb,
     yuv_region_2_bgr,
+    yuv_region_2_yuv,
 )
 
 logger = logging.getLogger(__name__)
@@ -97,7 +97,7 @@ def create_tensor_input(frame, model_config, region):
     elif model_config.input_pixel_format == PixelFormatEnum.bgr:
         cropped_frame = yuv_region_2_bgr(frame, region)
     else:
-        cropped_frame = yuv_crop_and_resize(frame, region)
+        cropped_frame = yuv_region_2_yuv(frame, region)
 
     # Resize if needed
     if cropped_frame.shape != (model_config.height, model_config.width, 3):