frigate/frigate/test/test_video.py

import unittest

import cv2
import numpy as np
from norfair.drawing.color import Palette
from norfair.drawing.drawer import Drawer

from frigate.video import (
    get_cluster_boundary,
    get_cluster_candidates,
    get_cluster_region,
)
from frigate.util import intersection


def draw_box(frame, box, color=(255, 0, 0), thickness=2):
    cv2.rectangle(
        frame,
        (box[0], box[1]),
        (box[2], box[3]),
        color,
        thickness,
    )


def save_clusters_image(name, boxes, candidates, regions=[]):
    canvas = np.zeros((1000, 2000, 3), np.uint8)
    for cluster in candidates:
        color = Palette.choose_color(np.random.rand())
        for b in cluster:
            box = boxes[b]
            draw_box(canvas, box, color, 2)
            # bottom right
            text_anchor = (
                box[2],
                box[3],
            )
            canvas = Drawer.text(
                canvas,
                str(b),
                position=text_anchor,
                size=None,
                color=(255, 255, 255),
                thickness=None,
            )
    for r in regions:
        draw_box(canvas, r, (0, 255, 0), 2)
    cv2.imwrite(
        f"debug/frames/{name}.jpg",
        canvas,
    )


def save_cluster_boundary_image(name, boxes, bounding_boxes):
    canvas = np.zeros((1000, 2000, 3), np.uint8)
    color = Palette.choose_color(np.random.rand())
    for box in boxes:
        draw_box(canvas, box, color, 2)
    for bound in bounding_boxes:
        draw_box(canvas, bound, (0, 255, 0), 2)
    cv2.imwrite(
        f"debug/frames/{name}.jpg",
        canvas,
    )


class TestRegion(unittest.TestCase):
    def setUp(self):
        self.frame_shape = (1000, 2000)
        self.min_region_size = 160

    def test_cluster_candidates(self):
        boxes = [(100, 100, 200, 200), (202, 150, 252, 200), (900, 900, 950, 950)]

        cluster_candidates = get_cluster_candidates(
            self.frame_shape, self.min_region_size, boxes
        )

        # save_clusters_image("cluster_candidates", boxes, cluster_candidates)

        assert len(cluster_candidates) == 2

    def test_cluster_boundary(self):
        boxes = [(100, 100, 200, 200), (215, 215, 325, 325)]
        boundary_boxes = [
            get_cluster_boundary(box, self.min_region_size) for box in boxes
        ]

        # save_cluster_boundary_image("bound", boxes, boundary_boxes)
        assert len(boundary_boxes) == 2

    def test_cluster_regions(self):
        boxes = [(100, 100, 200, 200), (202, 150, 252, 200), (900, 900, 950, 950)]

        cluster_candidates = get_cluster_candidates(
            self.frame_shape, self.min_region_size, boxes
        )

        regions = [
            get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
            for candidate in cluster_candidates
        ]

        # save_clusters_image("regions", boxes, cluster_candidates, regions)
        assert len(regions) == 2

    def test_box_too_small_for_cluster(self):
        boxes = [(100, 100, 600, 600), (655, 100, 700, 145)]

        cluster_candidates = get_cluster_candidates(
            self.frame_shape, self.min_region_size, boxes
        )

        regions = [
            get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
            for candidate in cluster_candidates
        ]

        save_clusters_image("too_small", boxes, cluster_candidates, regions)

        assert len(cluster_candidates) == 2
        assert len(regions) == 2

    def test_redundant_clusters(self):
        boxes = [(100, 100, 200, 200), (305, 305, 415, 415)]

        cluster_candidates = get_cluster_candidates(
            self.frame_shape, self.min_region_size, boxes
        )

        regions = [
            get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
            for candidate in cluster_candidates
        ]

        # save_clusters_image("redundant", boxes, cluster_candidates, regions)

        assert len(cluster_candidates) == 2
        assert all([len(c) == 1 for c in cluster_candidates])
        assert len(regions) == 2

    def test_combine_boxes(self):
        boxes = [
            (460, 0, 561, 144),
            (565, 0, 586, 71),
        ]

        # boundary_boxes = [get_cluster_boundary(box) for box in boxes]
        # save_cluster_boundary_image("combine_bound", boxes, boundary_boxes)

        cluster_candidates = get_cluster_candidates(
            self.frame_shape, self.min_region_size, boxes
        )

        regions = [
            get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
            for candidate in cluster_candidates
        ]

        # save_clusters_image("combine", boxes, cluster_candidates, regions)
        assert len(regions) == 1

    def test_dont_combine_boxes(self):
        boxes = [(460, 0, 532, 129), (586, 0, 606, 46)]

        # boundary_boxes = [get_cluster_boundary(box) for box in boxes]
        # save_cluster_boundary_image("dont_combine_bound", boxes, boundary_boxes)

        cluster_candidates = get_cluster_candidates(
            self.frame_shape, self.min_region_size, boxes
        )

        regions = [
            get_cluster_region(self.frame_shape, self.min_region_size, candidate, boxes)
            for candidate in cluster_candidates
        ]

        # save_clusters_image("dont_combine", boxes, cluster_candidates, regions)
        assert len(regions) == 2


class TestObjectBoundingBoxes(unittest.TestCase):
    def setUp(self) -> None:
        pass

    def test_box_intersection(self):
        box_a = [2012, 191, 2031, 205]
        box_b = [887, 92, 985, 151]
        box_c = [899, 128, 1080, 175]

        assert intersection(box_a, box_b) == None
        assert intersection(box_b, box_c) == [899, 128, 985, 151]

    def test_consolidate_objects(self):
        selected_objects = [
            (
                "car",  # valid detection, should not be consolidated
                0.82421875,
                (2044, 240, 2134, 297),
                5130,
                1.5789473684210527,
                (2008, 188, 2168, 348),
            ),
            (
                "car",  # valid detection, should not be consolidated
                0.76953125,
                (124, 239, 967, 719),
                404640,
                1.75625,
                (0, 0, 2968, 2968),
            ),
            (
                "car",  # valid detection, should not be consolidated
                0.71484375,
                (890, 93, 976, 150),
                4902,
                1.5087719298245614,
                (854, 40, 1014, 200),
            ),
            (
                "car",  # valid detection, should not be consolidated
                0.6640625,
                (1959, 188, 2005, 222),
                1564,
                1.3529411764705883,
                (1768, 0, 2560, 792),
            ),
            (
                "car",  # invalid detection, should be consolidated
                0.65234375,
                (2033, 193, 2048, 203),
                150,
                1.5,
                (2008, 188, 2168, 348),
            ),
        ]

        [
            (
                "car",
                0.82421875,
                (2044, 240, 2134, 297),
                5130,
                1.5789473684210527,
                (2008, 188, 2168, 348),
            ),
            (
                "car",
                0.76953125,
                (124, 239, 967, 719),
                404640,
                1.75625,
                (0, 0, 2968, 2968),
            ),
        ]