31 lines
964 B
Python
31 lines
964 B
Python
import dlib
|
|
import cv2 as cv
|
|
|
|
|
|
class Tracker:
|
|
def __init__(self):
|
|
self.tracker = dlib.correlation_tracker()
|
|
self.is_tracking = False
|
|
|
|
def start_tracking(self, img, bbox):
|
|
"""Initialize the tracker with the given bounding box."""
|
|
x, y, w, h = bbox
|
|
self.tracker.start_track(img, dlib.rectangle(x, y, x + w, y + h))
|
|
self.is_tracking = True
|
|
|
|
def update(self, img):
|
|
"""Update the tracker with the new frame."""
|
|
if not self.is_tracking:
|
|
return None
|
|
|
|
self.tracker.update(img)
|
|
pos = self.tracker.get_position()
|
|
x1, y1, x2, y2 = int(pos.left()), int(pos.top()), int(pos.right()), int(pos.bottom())
|
|
|
|
# Validate the bounding box
|
|
if (x2 - x1) > 0 and (y2 - y1) > 0:
|
|
return (x1, y1, x2 - x1, y2 - y1) # Valid bounding box
|
|
else:
|
|
self.is_tracking = False # Stop tracking if the box is invalid
|
|
return None
|