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
This commit is contained in:
Nicolas Mowen 2024-11-15 09:54:59 -07:00 committed by GitHub
parent 4eea541352
commit 7fdf42a56f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 58 additions and 23 deletions

View File

@ -4,6 +4,7 @@ import base64
import logging import logging
from collections import defaultdict from collections import defaultdict
from statistics import median from statistics import median
from typing import Optional
import cv2 import cv2
import numpy as np import numpy as np
@ -423,10 +424,11 @@ class TrackedObjectAttribute:
"box": self.box, "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.""" """Find the best attribute for each object and return its ID."""
best_object_area = None best_object_area = None
best_object_id = None best_object_id = None
best_object_label = None
for obj in objects: for obj in objects:
if not box_inside(obj["box"], self.box): if not box_inside(obj["box"], self.box):
@ -440,8 +442,15 @@ class TrackedObjectAttribute:
if best_object_area is None: if best_object_area is None:
best_object_area = object_area best_object_area = object_area
best_object_id = obj["id"] best_object_id = obj["id"]
elif object_area < best_object_area: best_object_label = obj["label"]
best_object_area = object_area else:
best_object_id = obj["id"] 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 return best_object_id

View File

@ -3,8 +3,10 @@
import datetime import datetime
import logging import logging
import subprocess as sp import subprocess as sp
import threading
from abc import ABC, abstractmethod 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 string import printable
from typing import AnyStr, Optional from typing import AnyStr, Optional
@ -731,32 +733,56 @@ class FrameManager(ABC):
pass pass
class DictFrameManager(FrameManager): class SharedMemory(_mpshm.SharedMemory):
def __init__(self): # https://github.com/python/cpython/issues/82300#issuecomment-2169035092
self.frames = {}
def create(self, name, size) -> AnyStr: __lock = threading.Lock()
mem = bytearray(size)
self.frames[name] = mem
return mem
def get(self, name, shape): def __init__(
mem = self.frames[name] self,
return np.ndarray(shape, dtype=np.uint8, buffer=mem) name: Optional[str] = None,
create: bool = False,
size: int = 0,
*,
track: bool = True,
) -> None:
self._track = track
def close(self, name): # if tracking, normal init will suffice
pass if track:
return super().__init__(name=name, create=create, size=size)
def delete(self, name): # lock so that other threads don't attempt to use the
del self.frames[name] # 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): class SharedMemoryFrameManager(FrameManager):
def __init__(self): def __init__(self):
self.shm_store: dict[str, shared_memory.SharedMemory] = {} self.shm_store: dict[str, SharedMemory] = {}
def create(self, name: str, size) -> AnyStr: 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 self.shm_store[name] = shm
return shm.buf return shm.buf
@ -765,7 +791,7 @@ class SharedMemoryFrameManager(FrameManager):
if name in self.shm_store: if name in self.shm_store:
shm = self.shm_store[name] shm = self.shm_store[name]
else: else:
shm = shared_memory.SharedMemory(name=name) shm = SharedMemory(name=name, track=False)
self.shm_store[name] = shm self.shm_store[name] = shm
return np.ndarray(shape, dtype=np.uint8, buffer=shm.buf) return np.ndarray(shape, dtype=np.uint8, buffer=shm.buf)
except FileNotFoundError: except FileNotFoundError:
@ -788,7 +814,7 @@ class SharedMemoryFrameManager(FrameManager):
del self.shm_store[name] del self.shm_store[name]
else: else:
try: try:
shm = shared_memory.SharedMemory(name=name) shm = SharedMemory(name=name, track=False)
shm.close() shm.close()
shm.unlink() shm.unlink()
except FileNotFoundError: except FileNotFoundError: