You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
146 lines
4.7 KiB
146 lines
4.7 KiB
import os
|
|
import sys
|
|
import cv2
|
|
import time # Import time module for adding delay
|
|
from icecream import ic
|
|
|
|
env_path = os.path.join(os.path.dirname(__file__), '')
|
|
if env_path not in sys.path:
|
|
sys.path.append(env_path)
|
|
|
|
from pytracking.evaluation import Tracker
|
|
|
|
|
|
def run_webcam(debug=None, refresh_rate=30):
|
|
tracker = Tracker('dimp', 'dimp50', debug)
|
|
|
|
class UIControl:
|
|
def __init__(self):
|
|
self.mode = 'init' # init, select, track
|
|
self.target_tl = (-1, -1)
|
|
self.target_br = (-1, -1)
|
|
self.new_init = False
|
|
|
|
def mouse_callback(self, event, x, y, flags, param):
|
|
if event == cv2.EVENT_LBUTTONDOWN and self.mode == 'init':
|
|
self.target_tl = (x, y)
|
|
self.target_br = (x, y)
|
|
self.mode = 'select'
|
|
elif event == cv2.EVENT_MOUSEMOVE and self.mode == 'select':
|
|
self.target_br = (x, y)
|
|
elif event == cv2.EVENT_LBUTTONDOWN and self.mode == 'select':
|
|
self.target_br = (x, y)
|
|
self.mode = 'init'
|
|
self.new_init = True
|
|
|
|
def get_tl(self):
|
|
return self.target_tl if self.target_tl[0] < self.target_br[0] else self.target_br
|
|
|
|
def get_br(self):
|
|
return self.target_br if self.target_tl[0] < self.target_br[0] else self.target_tl
|
|
|
|
def get_bb(self):
|
|
tl = self.get_tl()
|
|
br = self.get_br()
|
|
bb = [min(tl[0], br[0]), min(tl[1], br[1]), abs(br[0] - tl[0]), abs(br[1] - tl[1])]
|
|
return bb
|
|
|
|
ui_control = UIControl()
|
|
cap = cv2.VideoCapture(0)
|
|
#cap = cv2.VideoCapture('/home/mht/Downloads/car4.mp4')
|
|
#cap = cv2.VideoCapture('/home/mht/Downloads/messi.mp4')
|
|
display_name = 'tracker'
|
|
cv2.namedWindow(display_name, cv2.WINDOW_NORMAL)
|
|
cv2.setMouseCallback(display_name, ui_control.mouse_callback)
|
|
|
|
# Read the first frame before the loop starts
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
print("Error: Cannot read video.")
|
|
cap.release()
|
|
sys.exit()
|
|
|
|
frame_disp = frame.copy()
|
|
|
|
|
|
# Show the first frame and let the user draw a bounding box
|
|
while not ui_control.new_init:
|
|
|
|
temp_disp = frame_disp.copy()
|
|
if ui_control.mode == 'select':
|
|
cv2.rectangle(temp_disp, ui_control.get_tl(), ui_control.get_br(), (0, 255, 0), 2)
|
|
cv2.putText(temp_disp, "Draw bounding box and click to confirm...",
|
|
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
|
|
cv2.imshow(display_name, temp_disp)
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'): # Quit
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
sys.exit()
|
|
|
|
# Initialize the tracker with the bounding box selected
|
|
init_bbox = ui_control.get_bb()
|
|
ic(init_bbox)
|
|
tracker.init(frame, init_bbox)
|
|
|
|
#tracker.init(frame, [243, 159, 17, 14]) #fixed starting bb
|
|
|
|
# Frame delay for refresh rate
|
|
frame_delay = 1.0 / refresh_rate
|
|
|
|
# Main loop starts
|
|
while True:
|
|
|
|
start_time = time.time()
|
|
ret, frame = cap.read()
|
|
if not ret:
|
|
break # Exit if no more frames
|
|
frame_disp = frame.copy()
|
|
|
|
bbox, success = tracker.update(frame)
|
|
|
|
if bbox is not None:
|
|
color = (0, 255, 0) if success else (0, 0, 255)
|
|
cv2.rectangle(frame_disp, (bbox[0], bbox[1]), (bbox[0] + bbox[2], bbox[1] + bbox[3]),
|
|
color, 3)
|
|
|
|
cv2.imshow(display_name, frame_disp)
|
|
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'): # Quit
|
|
break
|
|
elif key == ord('r'): # Reset tracker
|
|
tracker.stop()
|
|
ui_control.mode = 'init'
|
|
ui_control.new_init = False
|
|
|
|
# Reinitialize on current frame
|
|
while not ui_control.new_init:
|
|
temp_disp = frame.copy()
|
|
if ui_control.mode == 'select':
|
|
cv2.rectangle(temp_disp, ui_control.get_tl(), ui_control.get_br(), (0, 255, 0), 2)
|
|
cv2.putText(temp_disp, "Draw bounding box and click to confirm...",
|
|
(10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
|
|
cv2.imshow(display_name, temp_disp)
|
|
key = cv2.waitKey(1)
|
|
if key == ord('q'): # Quit
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
sys.exit()
|
|
init_bbox = ui_control.get_bb()
|
|
tracker.init(frame, init_bbox)
|
|
|
|
# Maintain refresh rate
|
|
elapsed_time = time.time() - start_time
|
|
time.sleep(max(0, frame_delay - elapsed_time))
|
|
|
|
cap.release()
|
|
cv2.destroyAllWindows()
|
|
|
|
|
|
def main():
|
|
run_webcam(debug=False, refresh_rate=24)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|