allow region to extend beyond the frame

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():

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)

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
-
-    x1 = region[0] 
-    y1 = region[1]
-
-    uv_x1 = x1//2
-    uv_y1 = y1//4
-
-    uv_width = size//2
-    uv_height = size//4
-
-    u_y_start = height
-    v_y_start = height + height//4
-    two_x_offset = width//2
-
-    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]
-
-    return cv2.cvtColor(yuv_cropped_frame, cv2.COLOR_YUV2RGB_I420)
+    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]
+            ]
+
+        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 (