mirror of
https://github.com/blakeblackshear/frigate.git
synced 2024-11-21 19:07:46 +01:00
allow region to extend beyond the frame
This commit is contained in:
Blake Blackshear
parent
8b9c8a2e80
commit
41dd4447cc
@ -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():
|
||||
|
||||
39
frigate/test/test_yuv_region_2_rgb.py
Normal file
39
frigate/test/test_yuv_region_2_rgb.py
Normal file
@ -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)
|
||||
165
frigate/util.py
165
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):
|
||||
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
|
||||
|
||||
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]
|
||||
# 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 (
|
||||
|
||||
Loading…
Reference in New Issue
Block a user