From f4a0ec43a68a54210ebb40afc8c588ab4c91dd58 Mon Sep 17 00:00:00 2001 From: Blake Blackshear Date: Sat, 8 May 2021 08:27:27 -0500 Subject: [PATCH] wip yuv_rgb refactor --- frigate/util.py | 163 +++++++++++++++++++++++++----------------------- 1 file changed, 86 insertions(+), 77 deletions(-) diff --git a/frigate/util.py b/frigate/util.py index 17da6ca23..3d8650e97 100755 --- a/frigate/util.py +++ b/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}")