diff --git a/frigate/objects.py b/frigate/objects.py index c147d4653..f4cb6644e 100644 --- a/frigate/objects.py +++ b/frigate/objects.py @@ -12,7 +12,7 @@ import cv2 import numpy as np from scipy.spatial import distance as dist -from frigate.util import calculate_region, draw_box_with_label +from frigate.util import draw_box_with_label class ObjectTracker(): diff --git a/frigate/test/test_yuv_region_2_rgb.py b/frigate/test/test_yuv_region_2_rgb.py new file mode 100644 index 000000000..d8378fe79 --- /dev/null +++ b/frigate/test/test_yuv_region_2_rgb.py @@ -0,0 +1,39 @@ +import cv2 +import numpy as np +from unittest import TestCase, main +from frigate.util import yuv_region_2_rgb + +class TestYuvRegion2RGB(TestCase): + def setUp(self): + self.bgr_frame = np.zeros((100, 200, 3), np.uint8) + self.bgr_frame[:] = (0, 0, 255) + self.bgr_frame[5:55, 5:55] = (255,0,0) + # cv2.imwrite(f"bgr_frame.jpg", self.bgr_frame) + self.yuv_frame = cv2.cvtColor(self.bgr_frame, cv2.COLOR_BGR2YUV_I420) + + def test_crop_yuv(self): + cropped = yuv_region_2_rgb(self.yuv_frame, (10,10,50,50)) + # ensure the upper left pixel is blue + assert(np.all(cropped[0, 0] == [0, 0, 255])) + + def test_crop_yuv_out_of_bounds(self): + cropped = yuv_region_2_rgb(self.yuv_frame, (0,0,200,200)) + # cv2.imwrite(f"cropped.jpg", cv2.cvtColor(cropped, cv2.COLOR_RGB2BGR)) + # ensure the upper left pixel is red + # the yuv conversion has some noise + assert(np.all(cropped[0, 0] == [255, 1, 0])) + # ensure the bottom right is black + assert(np.all(cropped[199, 199] == [0, 0, 0])) + + def test_crop_yuv_portrait(self): + bgr_frame = np.zeros((1920, 1080, 3), np.uint8) + bgr_frame[:] = (0, 0, 255) + bgr_frame[5:55, 5:55] = (255,0,0) + # cv2.imwrite(f"bgr_frame.jpg", self.bgr_frame) + yuv_frame = cv2.cvtColor(bgr_frame, cv2.COLOR_BGR2YUV_I420) + + cropped = yuv_region_2_rgb(yuv_frame, (0, 852, 650, 1502)) + # cv2.imwrite(f"cropped.jpg", cv2.cvtColor(cropped, cv2.COLOR_RGB2BGR)) + +if __name__ == '__main__': + main(verbosity=2) \ No newline at end of file diff --git a/frigate/util.py b/frigate/util.py index 504bf885e..e44d4e9a8 100755 --- a/frigate/util.py +++ b/frigate/util.py @@ -47,14 +47,11 @@ def draw_box_with_label(frame, x_min, y_min, x_max, y_max, label, info, thicknes cv2.putText(frame, display_text, (text_offset_x, text_offset_y + line_height - 3), font, fontScale=font_scale, color=(0, 0, 0), thickness=2) def calculate_region(frame_shape, xmin, ymin, xmax, ymax, multiplier=2): - # size is larger than longest edge - size = int(max(xmax-xmin, ymax-ymin)*multiplier) + # size is the longest edge and divisible by 4 + size = int(max(xmax-xmin, ymax-ymin)//4*4*multiplier) # dont go any smaller than 300 if size < 300: size = 300 - # if the size is too big to fit in the frame - if size > min(frame_shape[0], frame_shape[1]): - size = min(frame_shape[0], frame_shape[1]) # x_offset is midpoint of bounding box minus half the size x_offset = int((xmax-xmin)/2.0+xmin-size/2.0) @@ -62,48 +59,156 @@ def calculate_region(frame_shape, xmin, ymin, xmax, ymax, multiplier=2): if x_offset < 0: x_offset = 0 elif x_offset > (frame_shape[1]-size): - x_offset = (frame_shape[1]-size) + x_offset = max(0, (frame_shape[1]-size)) # y_offset is midpoint of bounding box minus half the size y_offset = int((ymax-ymin)/2.0+ymin-size/2.0) - # if outside the image + # # if outside the image if y_offset < 0: y_offset = 0 elif y_offset > (frame_shape[0]-size): - y_offset = (frame_shape[0]-size) + y_offset = max(0, (frame_shape[0]-size)) return (x_offset, y_offset, x_offset+size, y_offset+size) +def get_yuv_crop(frame_shape, crop): + # crop should be (x1,y1,x2,y2) + frame_height = frame_shape[0]//3*2 + frame_width = frame_shape[1] + + # compute the width/height of the uv channels + uv_width = frame_width//2 # width of the uv channels + uv_height = frame_height//4 # height of the uv channels + + # compute the offset for upper left corner of the uv channels + uv_x_offset = crop[0]//2 # x offset of the uv channels + uv_y_offset = crop[1]//4 # y offset of the uv channels + + # compute the width/height of the uv crops + uv_crop_width = (crop[2] - crop[0])//2 # width of the cropped uv channels + uv_crop_height = (crop[3] - crop[1])//4 # height of the cropped uv channels + + # ensure crop dimensions are multiples of 2 and 4 + y = ( + crop[0], + crop[1], + crop[0] + uv_crop_width*2, + crop[1] + uv_crop_height*4 + ) + + u1 = ( + 0 + uv_x_offset, + frame_height + uv_y_offset, + 0 + uv_x_offset + uv_crop_width, + frame_height + uv_y_offset + uv_crop_height + ) + + u2 = ( + uv_width + uv_x_offset, + frame_height + uv_y_offset, + uv_width + uv_x_offset + uv_crop_width, + frame_height + uv_y_offset + uv_crop_height + ) + + v1 = ( + 0 + uv_x_offset, + frame_height + uv_height + uv_y_offset, + 0 + uv_x_offset + uv_crop_width, + frame_height + uv_height + uv_y_offset + uv_crop_height + ) + + v2 = ( + uv_width + uv_x_offset, + frame_height + uv_height + uv_y_offset, + uv_width + uv_x_offset + uv_crop_width, + frame_height + uv_height + uv_y_offset + uv_crop_height + ) + + return y, u1, u2, v1, v2 + def yuv_region_2_rgb(frame, region): - height = frame.shape[0]//3*2 - width = frame.shape[1] - # make sure the size is a multiple of 4 - size = (region[3] - region[1])//4*4 + try: + height = frame.shape[0]//3*2 + width = frame.shape[1] - x1 = region[0] - y1 = region[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) - uv_x1 = x1//2 - uv_y1 = y1//4 + y, u1, u2, v1, v2 = get_yuv_crop(frame.shape, crop_box) - uv_width = size//2 - uv_height = size//4 + # 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])) - u_y_start = height - v_y_start = height + height//4 - two_x_offset = width//2 + uv_channel_x_offset = y_channel_x_offset//2 + uv_channel_y_offset = y_channel_y_offset//4 - yuv_cropped_frame = np.zeros((size+size//2, size), np.uint8) - # y channel - yuv_cropped_frame[0:size, 0:size] = frame[y1:y1+size, x1:x1+size] - # u channel - yuv_cropped_frame[size:size+uv_height, 0:uv_width] = frame[uv_y1+u_y_start:uv_y1+u_y_start+uv_height, uv_x1:uv_x1+uv_width] - yuv_cropped_frame[size:size+uv_height, uv_width:size] = frame[uv_y1+u_y_start:uv_y1+u_y_start+uv_height, uv_x1+two_x_offset:uv_x1+two_x_offset+uv_width] - # v channel - yuv_cropped_frame[size+uv_height:size+uv_height*2, 0:uv_width] = frame[uv_y1+v_y_start:uv_y1+v_y_start+uv_height, uv_x1:uv_x1+uv_width] - yuv_cropped_frame[size+uv_height:size+uv_height*2, uv_width:size] = frame[uv_y1+v_y_start:uv_y1+v_y_start+uv_height, uv_x1+two_x_offset:uv_x1+two_x_offset+uv_width] + # 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 - return cv2.cvtColor(yuv_cropped_frame, cv2.COLOR_YUV2RGB_I420) + # 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 cv2.cvtColor(yuv_cropped_frame, cv2.COLOR_YUV2RGB_I420) + except: + print(f"frame.shape: {frame.shape}") + print(f"region: {region}") + raise def intersection(box_a, box_b): return (