34 lines
1.2 KiB
Python
34 lines
1.2 KiB
Python
class Counter:
|
|
def __init__(self, line_orientation, metrics):
|
|
self.in_count = 0
|
|
self.out_count = 0
|
|
self.track_memory = {}
|
|
self.line_orientation = line_orientation
|
|
self.metrics = metrics
|
|
|
|
def update(self, tracks, line_position):
|
|
for track in tracks:
|
|
track_id = track.id
|
|
x, y, w, h = track.bbox
|
|
|
|
|
|
if self.line_orientation == 'horizontal':
|
|
center = int(y + h / 2)
|
|
elif self.line_orientation == 'vertical':
|
|
center = int(x + w / 2)
|
|
else:
|
|
raise NotImplementedError(f'Line orientation {self.line_orientation} is invalid!')
|
|
|
|
if track_id not in self.track_memory:
|
|
self.track_memory[track_id] = center
|
|
continue
|
|
|
|
prev = self.track_memory[track_id]
|
|
self.track_memory[track_id] = center
|
|
|
|
if prev < line_position <= center:
|
|
self.in_count += 1
|
|
self.metrics["people_in"].inc()
|
|
elif prev > line_position >= center:
|
|
self.out_count += 1
|
|
self.metrics["people_out"].inc()
|