<filename>crossing_detector.py
import cv2
import sys
import settings
import tempfile
import numpy as np
from matplotlib import pyplot as plt
from scipy.spatial import distance as dist
from imageai.Detection import VideoObjectDetection

from utils import send_image_telegram

class CrossingDetector():

    """Основной класс с реализацией логики отслеживания объектов

    Attributes:
        ident_counter (int): Идентификатор для отслеживаемого объекта
        max_disappeared (int): Максимальное количество исчезновений
        trackable_objects (list): Список с отслеживаемыми объектами класса TrackableObject
        y_border (int): Граница
    """

    def __init__(self, y_border, max_disappeared=50):
        """Конструктор класса CrossingDetector

        Args:
            y_border (int): Граница
            max_disappeared (int): Максимальное количество исчезновений
        """
        self.trackable_objects = []
        self.ident_counter = 0
        self.y_border = y_border
        self.max_disappeared = max_disappeared

    class TrackableObject():

        """Класс для отслеживаемых объектов на видео

        Attributes:
            ident (int): Идентификатор объекта
            centroids (list): Список положений центроидов отслеживаемого объекта
            checked (bool): Пересечение границы зафиксированно
            count_disappeared (int): Количество исчезновений на кадрах
        """

        def __init__(self, ident, centroid):
            """Конструктор класса TrackableObject

            Args:
                ident (int): Идентификатор объекта
                centroid (tuple): Первый зафиксированный центроид объекта
            """
            self.ident = ident
            self.centroids = [centroid]
            self.checked = False
            self.count_disappeared = 0

        def add_centroid(self, centroid):
            """Метод для добавления нового зафиксированного центроида

            Args:
                centroid (tuple): Новый зафиксированный центроид
            """
            self.centroids.append(centroid)

        def get_direction(self):
            """Метод для вычисления направления отслеживаемого объекта

            Returns:
                result (float): Разница между текущей координатой y и средним предыдущих
            """
            y_coordinates = [p[1] for p in self.centroids]
            result = 0

            if y_coordinates:
                result = y_coordinates[-1] - np.mean(y_coordinates[:-1])

            return result

    def register_object(self, centroid):
        """Метод для регистрации нового объекта

        Args:
            centroid (tuple): Центроид
        """
        self.ident_counter += 1
        centroid = np.array(centroid)
        self.trackable_objects.append(
            CrossingDetector.TrackableObject(
                self.ident_counter,
                centroid
            )
        )

    def actualize_trackable_objects(self, detected_objects):
        """Метод для обновления центроидов отслеживаемых объектов

        Args:
            detected_objects (list): Зафиксированные объекты на кадре
        """
        detected_centroids = np.array(
            [self.get_center(detected_object.get('box_points'))
             for detected_object in detected_objects]
        )
        current_objects_centroids = np.array(
            [to.centroids[-1] for to in self.trackable_objects]
        )
        dist_matrix = dist.cdist(current_objects_centroids, detected_centroids)

        rows = dist_matrix.min(axis=1).argsort()
        cols = dist_matrix.argmin(axis=1)[rows]

        used_rows = set()
        used_cols = set()

        for (row, col) in zip(rows, cols):
            if row in used_rows or col in used_cols:
                continue

            self.trackable_objects[row].add_centroid(detected_centroids[col])
            self.trackable_objects[row].count_disappeared = 0
            used_rows.add(row)
            used_cols.add(col)

        unused_rows = set(range(0, dist_matrix.shape[0])).difference(used_rows)
        unused_cols = set(range(0, dist_matrix.shape[1])).difference(used_cols)

        if dist_matrix.shape[0] >= dist_matrix.shape[1]:
            for row in unused_rows:
                self.incr_object_count_disappeared(row)
        else:
            for col in unused_cols:
                self.register_object(detected_centroids[col])

    def incr_object_count_disappeared(self, index):
        """Метод для увеличения счетчика исчезновений объекта

        Args:
            index (int): Индекс объекта в списке
        """
        self.trackable_objects[index].count_disappeared += 1
        if self.trackable_objects[index].count_disappeared > self.max_disappeared:
            del self.trackable_objects[index]

    def register_detected_objects(self, objects):
        """Метод для регистрации новых объектов

        Args:
            objects (list): Список новых зафиксированных объектов
        """
        [self.register_object(self.get_center(detected_object.get(
            'box_points'))) for detected_object in objects]

    def get_center(self, box_points):
        """Метод для вычисления координат центроида

        Args:
            box_points (tuple): Координаты (x1, y1, x2, y2) рамки зафиксированного объекта

        Returns:
            tuple: Координаты центра рамки
        """
        x1, y1, x2, y2 = box_points
        def asix_center(v1, v2): return int((v2 - v1) / 2) + v1
        return asix_center(x1, x2), asix_center(y1, y2)

    def process_and_send_image(self, centroid, frame_img):
        """Метод для предобработки и отправки изображения

        Args:
            centroid (tuple): Центроид
            frame_img (array): Изображение
        """
        image = cv2.cvtColor(frame_img, cv2.COLOR_BGR2RGB)
        image = np.array(image)
        shape = image.shape
        plt.imshow(image, aspect='auto')
        plt.scatter(centroid[0], centroid[1], s=30, color='blue')
        plt.hlines(y=self.y_border, xmin=0, xmax=(shape[1] - 1), color='green')

        with tempfile.NamedTemporaryFile(suffix=".png") as tmpfile:
            plt.savefig(tmpfile, format="png")
            send_image_telegram(tmpfile.name)

    def process_frame(self, frame):
        """Метод для обработки кадра

        Args:
            frame (dict): Кадр
        """
        frame_num, objects, count_objects, frame_img = frame

        if frame_num == 1:
            self.register_detected_objects(objects)
            return

        if not objects:
            for i, to in enumerate(self.trackable_objects):
                self.incr_object_count_disappeared(i)
            return

        self.actualize_trackable_objects(objects)

        for to in self.trackable_objects:
            if not to.checked:
                to_direction = to.get_direction()
                curr_centroid = to.centroids[-1]
                first_centroid = to.centroids[0]

                if (to_direction > 0 and curr_centroid[1] > self.y_border and first_centroid[1] < self.y_border) or (
                        to_direction < 0 and curr_centroid[1] < self.y_border and first_centroid[1] > self.y_border):
                    self.process_and_send_image(curr_centroid, frame_img)
                    to.checked = True

<filename>crossing_detector.py
import cv2
import sys
import settings
import tempfile
import numpy as np
from matplotlib import pyplot as plt
from scipy.spatial import distance as dist
from imageai.Detection import VideoObjectDetection

from utils import send_image_telegram

class CrossingDetector():

    """Основной класс с реализацией логики отслеживания объектов

    Attributes:
        ident_counter (int): Идентификатор для отслеживаемого объекта
        max_disappeared (int): Максимальное количество исчезновений
        trackable_objects (list): Список с отслеживаемыми объектами класса TrackableObject
        y_border (int): Граница
    """

    def __init__(self, y_border, max_disappeared=50):
        """Конструктор класса CrossingDetector

        Args:
            y_border (int): Граница
            max_disappeared (int): Максимальное количество исчезновений
        """
        self.trackable_objects = []
        self.ident_counter = 0
        self.y_border = y_border
        self.max_disappeared = max_disappeared

    class TrackableObject():

        """Класс для отслеживаемых объектов на видео

        Attributes:
            ident (int): Идентификатор объекта
            centroids (list): Список положений центроидов отслеживаемого объекта
            checked (bool): Пересечение границы зафиксированно
            count_disappeared (int): Количество исчезновений на кадрах
        """

        def __init__(self, ident, centroid):
            """Конструктор класса TrackableObject

            Args:
                ident (int): Идентификатор объекта
                centroid (tuple): Первый зафиксированный центроид объекта
            """
            self.ident = ident
            self.centroids = [centroid]
            self.checked = False
            self.count_disappeared = 0

        def add_centroid(self, centroid):
            """Метод для добавления нового зафиксированного центроида

            Args:
                centroid (tuple): Новый зафиксированный центроид
            """
            self.centroids.append(centroid)

        def get_direction(self):
            """Метод для вычисления направления отслеживаемого объекта

            Returns:
                result (float): Разница между текущей координатой y и средним предыдущих
            """
            y_coordinates = [p[1] for p in self.centroids]
            result = 0

            if y_coordinates:
                result = y_coordinates[-1] - np.mean(y_coordinates[:-1])

            return result

    def register_object(self, centroid):
        """Метод для регистрации нового объекта

        Args:
            centroid (tuple): Центроид
        """
        self.ident_counter += 1
        centroid = np.array(centroid)
        self.trackable_objects.append(
            CrossingDetector.TrackableObject(
                self.ident_counter,
                centroid
            )
        )

    def actualize_trackable_objects(self, detected_objects):
        """Метод для обновления центроидов отслеживаемых объектов

        Args:
            detected_objects (list): Зафиксированные объекты на кадре
        """
        detected_centroids = np.array(
            [self.get_center(detected_object.get('box_points'))
             for detected_object in detected_objects]
        )
        current_objects_centroids = np.array(
            [to.centroids[-1] for to in self.trackable_objects]
        )
        dist_matrix = dist.cdist(current_objects_centroids, detected_centroids)

        rows = dist_matrix.min(axis=1).argsort()
        cols = dist_matrix.argmin(axis=1)[rows]

        used_rows = set()
        used_cols = set()

        for (row, col) in zip(rows, cols):
            if row in used_rows or col in used_cols:
                continue

            self.trackable_objects[row].add_centroid(detected_centroids[col])
            self.trackable_objects[row].count_disappeared = 0
            used_rows.add(row)
            used_cols.add(col)

        unused_rows = set(range(0, dist_matrix.shape[0])).difference(used_rows)
        unused_cols = set(range(0, dist_matrix.shape[1])).difference(used_cols)

        if dist_matrix.shape[0] >= dist_matrix.shape[1]:
            for row in unused_rows:
                self.incr_object_count_disappeared(row)
        else:
            for col in unused_cols:
                self.register_object(detected_centroids[col])

    def incr_object_count_disappeared(self, index):
        """Метод для увеличения счетчика исчезновений объекта

        Args:
            index (int): Индекс объекта в списке
        """
        self.trackable_objects[index].count_disappeared += 1
        if self.trackable_objects[index].count_disappeared > self.max_disappeared:
            del self.trackable_objects[index]

    def register_detected_objects(self, objects):
        """Метод для регистрации новых объектов

        Args:
            objects (list): Список новых зафиксированных объектов
        """
        [self.register_object(self.get_center(detected_object.get(
            'box_points'))) for detected_object in objects]

    def get_center(self, box_points):
        """Метод для вычисления координат центроида

        Args:
            box_points (tuple): Координаты (x1, y1, x2, y2) рамки зафиксированного объекта

        Returns:
            tuple: Координаты центра рамки
        """
        x1, y1, x2, y2 = box_points
        def asix_center(v1, v2): return int((v2 - v1) / 2) + v1
        return asix_center(x1, x2), asix_center(y1, y2)

    def process_and_send_image(self, centroid, frame_img):
        """Метод для предобработки и отправки изображения

        Args:
            centroid (tuple): Центроид
            frame_img (array): Изображение
        """
        image = cv2.cvtColor(frame_img, cv2.COLOR_BGR2RGB)
        image = np.array(image)
        shape = image.shape
        plt.imshow(image, aspect='auto')
        plt.scatter(centroid[0], centroid[1], s=30, color='blue')
        plt.hlines(y=self.y_border, xmin=0, xmax=(shape[1] - 1), color='green')

        with tempfile.NamedTemporaryFile(suffix=".png") as tmpfile:
            plt.savefig(tmpfile, format="png")
            send_image_telegram(tmpfile.name)

    def process_frame(self, frame):
        """Метод для обработки кадра

        Args:
            frame (dict): Кадр
        """
        frame_num, objects, count_objects, frame_img = frame

        if frame_num == 1:
            self.register_detected_objects(objects)
            return

        if not objects:
            for i, to in enumerate(self.trackable_objects):
                self.incr_object_count_disappeared(i)
            return

        self.actualize_trackable_objects(objects)

        for to in self.trackable_objects:
            if not to.checked:
                to_direction = to.get_direction()
                curr_centroid = to.centroids[-1]
                first_centroid = to.centroids[0]

                if (to_direction > 0 and curr_centroid[1] > self.y_border and first_centroid[1] < self.y_border) or (
                        to_direction < 0 and curr_centroid[1] < self.y_border and first_centroid[1] > self.y_border):
                    self.process_and_send_image(curr_centroid, frame_img)
                    to.checked = True

Основной класс с реализацией логики отслеживания объектов      Attributes:         ident_counter (int): Идентификатор для отслеживаемого объекта         max_disappeared (int): Максимальное количество исчезновений         trackable_objects (list): Список с отслеживаемыми объектами класса TrackableObject         y_border (int): Граница Конструктор класса CrossingDetector          Args:             y_border (int): Граница             max_disappeared (int): Максимальное количество исчезновений Класс для отслеживаемых объектов на видео          Attributes:             ident (int): Идентификатор объекта             centroids (list): Список положений центроидов отслеживаемого объекта             checked (bool): Пересечение границы зафиксированно             count_disappeared (int): Количество исчезновений на кадрах Конструктор класса TrackableObject              Args:                 ident (int): Идентификатор объекта                 centroid (tuple): Первый зафиксированный центроид объекта Метод для добавления нового зафиксированного центроида              Args:                 centroid (tuple): Новый зафиксированный центроид Метод для вычисления направления отслеживаемого объекта              Returns:                 result (float): Разница между текущей координатой y и средним предыдущих Метод для регистрации нового объекта          Args:             centroid (tuple): Центроид Метод для обновления центроидов отслеживаемых объектов          Args:             detected_objects (list): Зафиксированные объекты на кадре Метод для увеличения счетчика исчезновений объекта          Args:             index (int): Индекс объекта в списке Метод для регистрации новых объектов          Args:             objects (list): Список новых зафиксированных объектов Метод для вычисления координат центроида          Args:             box_points (tuple): Координаты (x1, y1, x2, y2) рамки зафиксированного объекта          Returns:             tuple: Координаты центра рамки Метод для предобработки и отправки изображения          Args:             centroid (tuple): Центроид             frame_img (array): Изображение Метод для обработки кадра          Args:             frame (dict): Кадр