wip yuv_rgb refactor

frigate/util.py

@@ -149,85 +149,94 @@ def get_yuv_crop(frame_shape, crop):
  return y, u1, u2, v1, v2
 
 
+def yuv_crop_and_resize(frame, region, height=None):
+ # Crops and resizes a YUV frame while maintaining aspect ratio
+ # https://stackoverflow.com/a/57022634
+ height = frame.shape[0] // 3 * 2
+ width = frame.shape[1]
+
+ # get the crop box if the region extends beyond the frame
+ crop_x1 = max(0, region[0])
+ crop_y1 = max(0, region[1])
+ # ensure these are a multiple of 4
+ crop_x2 = min(width, region[2])
+ crop_y2 = min(height, region[3])
+ crop_box = (crop_x1, crop_y1, crop_x2, crop_y2)
+
+ y, u1, u2, v1, v2 = get_yuv_crop(frame.shape, crop_box)
+
+ # if the region starts outside the frame, indent the start point in the cropped frame
+ y_channel_x_offset = abs(min(0, region[0]))
+ y_channel_y_offset = abs(min(0, region[1]))
+
+ uv_channel_x_offset = y_channel_x_offset // 2
+ uv_channel_y_offset = y_channel_y_offset // 4
+
+ # create the yuv region frame
+ # make sure the size is a multiple of 4
+ # TODO: this should be based on the size after resize now
+ size = (region[3] - region[1]) // 4 * 4
+ yuv_cropped_frame = np.zeros((size + size // 2, size), np.uint8)
+ # fill in black
+ yuv_cropped_frame[:] = 128
+ yuv_cropped_frame[0:size, 0:size] = 16
+
+ # copy the y channel
+ yuv_cropped_frame[
+ y_channel_y_offset : y_channel_y_offset + y[3] - y[1],
+ y_channel_x_offset : y_channel_x_offset + y[2] - y[0],
+ ] = frame[y[1] : y[3], y[0] : y[2]]
+
+ uv_crop_width = u1[2] - u1[0]
+ uv_crop_height = u1[3] - u1[1]
+
+ # copy u1
+ yuv_cropped_frame[
+ size + uv_channel_y_offset : size + uv_channel_y_offset + uv_crop_height,
+ 0 + uv_channel_x_offset : 0 + uv_channel_x_offset + uv_crop_width,
+ ] = frame[u1[1] : u1[3], u1[0] : u1[2]]
+
+ # copy u2
+ yuv_cropped_frame[
+ size + uv_channel_y_offset : size + uv_channel_y_offset + uv_crop_height,
+ size // 2
+ + uv_channel_x_offset : size // 2
+ + uv_channel_x_offset
+ + uv_crop_width,
+ ] = frame[u2[1] : u2[3], u2[0] : u2[2]]
+
+ # copy v1
+ yuv_cropped_frame[
+ size
+ + size // 4
+ + uv_channel_y_offset : size
+ + size // 4
+ + uv_channel_y_offset
+ + uv_crop_height,
+ 0 + uv_channel_x_offset : 0 + uv_channel_x_offset + uv_crop_width,
+ ] = frame[v1[1] : v1[3], v1[0] : v1[2]]
+
+ # copy v2
+ yuv_cropped_frame[
+ size
+ + size // 4
+ + uv_channel_y_offset : size
+ + size // 4
+ + uv_channel_y_offset
+ + uv_crop_height,
+ size // 2
+ + uv_channel_x_offset : size // 2
+ + uv_channel_x_offset
+ + uv_crop_width,
+ ] = frame[v2[1] : v2[3], v2[0] : v2[2]]
+
+ return yuv_cropped_frame
+
+
 def yuv_region_2_rgb(frame, region):
  try:
- height = frame.shape[0] // 3 * 2
- width = frame.shape[1]
-
- # get the crop box if the region extends beyond the frame
- crop_x1 = max(0, region[0])
- crop_y1 = max(0, region[1])
- # ensure these are a multiple of 4
- crop_x2 = min(width, region[2])
- crop_y2 = min(height, region[3])
- crop_box = (crop_x1, crop_y1, crop_x2, crop_y2)
-
- y, u1, u2, v1, v2 = get_yuv_crop(frame.shape, crop_box)
-
- # if the region starts outside the frame, indent the start point in the cropped frame
- y_channel_x_offset = abs(min(0, region[0]))
- y_channel_y_offset = abs(min(0, region[1]))
-
- uv_channel_x_offset = y_channel_x_offset // 2
- uv_channel_y_offset = y_channel_y_offset // 4
-
- # create the yuv region frame
- # make sure the size is a multiple of 4
- size = (region[3] - region[1]) // 4 * 4
- yuv_cropped_frame = np.zeros((size + size // 2, size), np.uint8)
- # fill in black
- yuv_cropped_frame[:] = 128
- yuv_cropped_frame[0:size, 0:size] = 16
-
- # copy the y channel
- yuv_cropped_frame[
- y_channel_y_offset : y_channel_y_offset + y[3] - y[1],
- y_channel_x_offset : y_channel_x_offset + y[2] - y[0],
- ] = frame[y[1] : y[3], y[0] : y[2]]
-
- uv_crop_width = u1[2] - u1[0]
- uv_crop_height = u1[3] - u1[1]
-
- # copy u1
- yuv_cropped_frame[
- size + uv_channel_y_offset : size + uv_channel_y_offset + uv_crop_height,
- 0 + uv_channel_x_offset : 0 + uv_channel_x_offset + uv_crop_width,
- ] = frame[u1[1] : u1[3], u1[0] : u1[2]]
-
- # copy u2
- yuv_cropped_frame[
- size + uv_channel_y_offset : size + uv_channel_y_offset + uv_crop_height,
- size // 2
- + uv_channel_x_offset : size // 2
- + uv_channel_x_offset
- + uv_crop_width,
- ] = frame[u2[1] : u2[3], u2[0] : u2[2]]
-
- # copy v1
- yuv_cropped_frame[
- size
- + size // 4
- + uv_channel_y_offset : size
- + size // 4
- + uv_channel_y_offset
- + uv_crop_height,
- 0 + uv_channel_x_offset : 0 + uv_channel_x_offset + uv_crop_width,
- ] = frame[v1[1] : v1[3], v1[0] : v1[2]]
-
- # copy v2
- yuv_cropped_frame[
- size
- + size // 4
- + uv_channel_y_offset : size
- + size // 4
- + uv_channel_y_offset
- + uv_crop_height,
- size // 2
- + uv_channel_x_offset : size // 2
- + uv_channel_x_offset
- + uv_crop_width,
- ] = frame[v2[1] : v2[3], v2[0] : v2[2]]
-
+ # TODO: does this copy the numpy array?
+ yuv_cropped_frame = yuv_crop_and_resize(frame, region)
  return cv2.cvtColor(yuv_cropped_frame, cv2.COLOR_YUV2RGB_I420)
  except:
  print(f"frame.shape: {frame.shape}")