From 7fdf42a56f452285ac296dcbaf4301c8ac154dbf Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Fri, 15 Nov 2024 09:54:59 -0700
Subject: [PATCH] Various Fixes (#15004)

* Don't track shared memory in frame tracker

* Don't track any instance

* Don't assign sub label to objects when multiple cars are overlapping

* Formatting

* Fix assignment
---
 frigate/track/tracked_object.py | 17 ++++++---
 frigate/util/image.py           | 64 +++++++++++++++++++++++----------
 2 files changed, 58 insertions(+), 23 deletions(-)

diff --git a/frigate/track/tracked_object.py b/frigate/track/tracked_object.py
index a4b4e8426..65e7a2ed5 100644
--- a/frigate/track/tracked_object.py
+++ b/frigate/track/tracked_object.py
@@ -4,6 +4,7 @@ import base64
 import logging
 from collections import defaultdict
 from statistics import median
+from typing import Optional
 
 import cv2
 import numpy as np
@@ -423,10 +424,11 @@ class TrackedObjectAttribute:
             "box": self.box,
         }
 
-    def find_best_object(self, objects: list[dict[str, any]]) -> str:
+    def find_best_object(self, objects: list[dict[str, any]]) -> Optional[str]:
         """Find the best attribute for each object and return its ID."""
         best_object_area = None
         best_object_id = None
+        best_object_label = None
 
         for obj in objects:
             if not box_inside(obj["box"], self.box):
@@ -440,8 +442,15 @@ class TrackedObjectAttribute:
             if best_object_area is None:
                 best_object_area = object_area
                 best_object_id = obj["id"]
-            elif object_area < best_object_area:
-                best_object_area = object_area
-                best_object_id = obj["id"]
+                best_object_label = obj["label"]
+            else:
+                if best_object_label == "car" and obj["label"] == "car":
+                    # if multiple cars are overlapping with the same label then the label will not be assigned
+                    return None
+                elif object_area < best_object_area:
+                    # if a car and person are overlapping then assign the label to the smaller object (which should be the person)
+                    best_object_area = object_area
+                    best_object_id = obj["id"]
+                    best_object_label = obj["label"]
 
         return best_object_id
diff --git a/frigate/util/image.py b/frigate/util/image.py
index 484737f71..763f0dfab 100644
--- a/frigate/util/image.py
+++ b/frigate/util/image.py
@@ -3,8 +3,10 @@
 import datetime
 import logging
 import subprocess as sp
+import threading
 from abc import ABC, abstractmethod
-from multiprocessing import shared_memory
+from multiprocessing import resource_tracker as _mprt
+from multiprocessing import shared_memory as _mpshm
 from string import printable
 from typing import AnyStr, Optional
 
@@ -731,32 +733,56 @@ class FrameManager(ABC):
         pass
 
 
-class DictFrameManager(FrameManager):
-    def __init__(self):
-        self.frames = {}
+class SharedMemory(_mpshm.SharedMemory):
+    # https://github.com/python/cpython/issues/82300#issuecomment-2169035092
 
-    def create(self, name, size) -> AnyStr:
-        mem = bytearray(size)
-        self.frames[name] = mem
-        return mem
+    __lock = threading.Lock()
 
-    def get(self, name, shape):
-        mem = self.frames[name]
-        return np.ndarray(shape, dtype=np.uint8, buffer=mem)
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        create: bool = False,
+        size: int = 0,
+        *,
+        track: bool = True,
+    ) -> None:
+        self._track = track
 
-    def close(self, name):
-        pass
+        # if tracking, normal init will suffice
+        if track:
+            return super().__init__(name=name, create=create, size=size)
 
-    def delete(self, name):
-        del self.frames[name]
+        # lock so that other threads don't attempt to use the
+        # register function during this time
+        with self.__lock:
+            # temporarily disable registration during initialization
+            orig_register = _mprt.register
+            _mprt.register = self.__tmp_register
+
+            # initialize; ensure original register function is
+            # re-instated
+            try:
+                super().__init__(name=name, create=create, size=size)
+            finally:
+                _mprt.register = orig_register
+
+    @staticmethod
+    def __tmp_register(*args, **kwargs) -> None:
+        return
+
+    def unlink(self) -> None:
+        if _mpshm._USE_POSIX and self._name:
+            _mpshm._posixshmem.shm_unlink(self._name)
+            if self._track:
+                _mprt.unregister(self._name, "shared_memory")
 
 
 class SharedMemoryFrameManager(FrameManager):
     def __init__(self):
-        self.shm_store: dict[str, shared_memory.SharedMemory] = {}
+        self.shm_store: dict[str, SharedMemory] = {}
 
     def create(self, name: str, size) -> AnyStr:
-        shm = shared_memory.SharedMemory(name=name, create=True, size=size)
+        shm = SharedMemory(name=name, create=True, size=size, track=False)
         self.shm_store[name] = shm
         return shm.buf
 
@@ -765,7 +791,7 @@ class SharedMemoryFrameManager(FrameManager):
             if name in self.shm_store:
                 shm = self.shm_store[name]
             else:
-                shm = shared_memory.SharedMemory(name=name)
+                shm = SharedMemory(name=name, track=False)
                 self.shm_store[name] = shm
             return np.ndarray(shape, dtype=np.uint8, buffer=shm.buf)
         except FileNotFoundError:
@@ -788,7 +814,7 @@ class SharedMemoryFrameManager(FrameManager):
             del self.shm_store[name]
         else:
             try:
-                shm = shared_memory.SharedMemory(name=name)
+                shm = SharedMemory(name=name, track=False)
                 shm.close()
                 shm.unlink()
             except FileNotFoundError: