Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Sign up
Appearance settings

Commit 13baf08

Browse files
committed
tracker work paralell with detector, work var.
chsnge req in ultralytics lib ( because need for tracker spec ver of opencv ) custom - for test
1 parent e1f9597 commit 13baf08

File tree

1 file changed

+367
-0
lines changed

1 file changed

+367
-0
lines changed

‎custom_paralel_threade.py

Lines changed: 367 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,367 @@
1+
#!/usr/bin/env python
2+
# -*- coding: utf-8 -*-
3+
import re
4+
import sys
5+
import copy
6+
import time
7+
import argparse
8+
import threading
9+
import queue
10+
import logging
11+
from collections import deque
12+
13+
import cv2 as cv
14+
import numpy as np
15+
from ultralytics import YOLO
16+
17+
# Configure logging
18+
logging.basicConfig(
19+
level=logging.DEBUG, # Set to DEBUG to capture all levels of logs
20+
format='[%(asctime)s] %(levelname)s - %(message)s',
21+
handlers=[
22+
logging.FileHandler("tracker_debug.log"), # Log to file
23+
logging.StreamHandler(sys.stdout) # Also log to console
24+
]
25+
)
26+
27+
def get_args():
28+
parser = argparse.ArgumentParser(description="Single Object Tracker with YOLOv8 Detection")
29+
30+
parser.add_argument("--device", default="sample_movie/bird.mp4",
31+
help="Video source. Can be a video file path or a camera index (integer).")
32+
parser.add_argument("--width", help='Capture width', type=int, default=960)
33+
parser.add_argument("--height", help='Capture height', type=int, default=540)
34+
35+
# Single tracker selection using mutually exclusive group
36+
tracker_group = parser.add_mutually_exclusive_group()
37+
tracker_group.add_argument('--mil', action='store_true', help='Use MIL tracker')
38+
tracker_group.add_argument('--goturn', action='store_true', help='Use GOTURN tracker')
39+
tracker_group.add_argument('--dasiamrpn', action='store_true', help='Use DaSiamRPN tracker')
40+
tracker_group.add_argument('--nano', action='store_true', help='Use Nano tracker')
41+
tracker_group.add_argument('--vit', action='store_true', help='Use ViT tracker')
42+
tracker_group.add_argument('--csrt', action='store_true', help='Use CSRT tracker')
43+
tracker_group.add_argument('--kcf', action='store_true', help='Use KCF tracker')
44+
tracker_group.add_argument('--boosting', action='store_true', help='Use Boosting tracker')
45+
tracker_group.add_argument('--mosse', action='store_true', help='Use MOSSE tracker')
46+
tracker_group.add_argument('--medianflow', action='store_true', help='Use MedianFlow tracker')
47+
tracker_group.add_argument('--tld', action='store_true', help='Use TLD tracker')
48+
49+
args = parser.parse_args()
50+
logging.debug(f"Parsed arguments: {args}")
51+
52+
return args
53+
54+
def isint(s):
55+
p = '[-+]?\d+'
56+
logging.debug(f"Checking if '{s}' is an integer.")
57+
return True if re.fullmatch(p, s) else False
58+
59+
def detect_objects(model, frame, detection_queue):
60+
"""
61+
Perform object detection and put the results in the detection_queue.
62+
"""
63+
try:
64+
results = model(frame)
65+
bboxes = []
66+
for result in results:
67+
for box in result.boxes:
68+
x1, y1, x2, y2 = box.xyxy[0] # Get the bounding box coordinates
69+
bboxes.append((int(x1), int(y1), int(x2 - x1), int(y2 - y1)))
70+
detection_queue.put(bboxes)
71+
logging.debug(f"Detected {len(bboxes)} bounding boxes.")
72+
except Exception as e:
73+
logging.error(f"Error during object detection: {e}")
74+
75+
class TrackerInfo:
76+
"""
77+
A class to hold a single tracker and its bounding box history for smoothing.
78+
"""
79+
def __init__(self, tracker, algorithm, initial_bbox, history_size=5):
80+
self.tracker = tracker
81+
self.algorithm = algorithm
82+
self.bbox_history = deque(maxlen=history_size)
83+
self.bbox_history.append(initial_bbox)
84+
self.color = None # To be assigned later
85+
86+
def update(self, frame):
87+
"""
88+
Update the tracker with the current frame and store bbox.
89+
Returns the smoothed bbox and tracking status.
90+
"""
91+
ok, bbox = self.tracker.update(frame)
92+
if ok:
93+
# Convert bbox to int and store
94+
bbox = [int(b) for b in bbox]
95+
self.bbox_history.append(bbox)
96+
# Calculate smoothed bbox
97+
smoothed_bbox = self.get_smoothed_bbox()
98+
return smoothed_bbox, True
99+
else:
100+
return None, False
101+
102+
def get_smoothed_bbox(self):
103+
"""
104+
Calculate the moving average of the bounding boxes.
105+
"""
106+
avg_bbox = np.mean(self.bbox_history, axis=0)
107+
return [int(coord) for coord in avg_bbox]
108+
109+
def initialize_tracker(image, tracker_algorithm, initial_bbox):
110+
"""
111+
Initialize a single tracker based on the specified algorithm and return TrackerInfo.
112+
"""
113+
tracker = None
114+
logging.debug(f"Initializing tracker '{tracker_algorithm}'.")
115+
116+
if tracker_algorithm == 'MIL':
117+
tracker = cv.TrackerMIL_create()
118+
elif tracker_algorithm == 'GOTURN':
119+
params = cv.TrackerGOTURN_Params()
120+
params.modelTxt = "model/GOTURN/goturn.prototxt"
121+
params.modelBin = "model/GOTURN/goturn.caffemodel"
122+
tracker = cv.TrackerGOTURN_create(params)
123+
elif tracker_algorithm == 'DaSiamRPN':
124+
params = cv.TrackerDaSiamRPN_Params()
125+
params.model = "model/DaSiamRPN/dasiamrpn_model.onnx"
126+
params.kernel_r1 = "model/DaSiamRPN/dasiamrpn_kernel_r1.onnx"
127+
params.kernel_cls1 = "model/DaSiamRPN/dasiamrpn_kernel_cls1.onnx"
128+
tracker = cv.TrackerDaSiamRPN_create(params)
129+
elif tracker_algorithm == 'Nano':
130+
params = cv.TrackerNano_Params()
131+
params.backbone = "model/nanotrackv2/nanotrack_backbone_sim.onnx"
132+
params.neckhead = "model/nanotrackv2/nanotrack_head_sim.onnx"
133+
tracker = cv.TrackerNano_create(params)
134+
elif tracker_algorithm == 'Vit':
135+
params = cv.TrackerVit_Params()
136+
params.net = "model/vit/object_tracking_vittrack_2023sep.onnx"
137+
tracker = cv.TrackerVit_create(params)
138+
elif tracker_algorithm == 'CSRT':
139+
tracker = cv.TrackerCSRT_create()
140+
elif tracker_algorithm == 'KCF':
141+
tracker = cv.TrackerKCF_create()
142+
elif tracker_algorithm == 'Boosting':
143+
tracker = cv.legacy_TrackerBoosting.create()
144+
elif tracker_algorithm == 'MOSSE':
145+
tracker = cv.legacy_TrackerMOSSE.create()
146+
elif tracker_algorithm == 'MedianFlow':
147+
tracker = cv.legacy_TrackerMedianFlow.create()
148+
elif tracker_algorithm == 'TLD':
149+
tracker = cv.legacy_TrackerTLD.create()
150+
else:
151+
logging.warning(f"Unknown tracker algorithm: {tracker_algorithm}")
152+
return None
153+
154+
if tracker is not None:
155+
try:
156+
tracker.init(image, tuple(initial_bbox))
157+
tracker_info = TrackerInfo(tracker, tracker_algorithm, initial_bbox)
158+
logging.debug(f"Successfully initialized '{tracker_algorithm}' tracker with bbox: {initial_bbox}")
159+
return tracker_info
160+
except Exception as e:
161+
logging.error(f"Exception during tracker initialization for '{tracker_algorithm}' with bbox {initial_bbox}: {e}")
162+
return None
163+
else:
164+
logging.error(f"Failed to initialize '{tracker_algorithm}' tracker with bbox: {initial_bbox}")
165+
return None
166+
167+
def main():
168+
color_list = [
169+
[255, 0, 0], # Blue
170+
[0, 255, 0], # Green
171+
[0, 0, 255], # Red
172+
[255, 255, 0], # Cyan
173+
[255, 0, 255], # Magenta
174+
[0, 255, 255], # Yellow
175+
[128, 0, 128], # Purple
176+
[128, 128, 0], # Olive
177+
[0, 128, 128], # Teal
178+
[128, 0, 0], # Maroon
179+
]
180+
181+
# Parse arguments ########################################################
182+
args = get_args()
183+
184+
cap_device = args.device
185+
cap_width = args.width
186+
cap_height = args.height
187+
188+
# Determine selected tracker algorithm
189+
tracker_algorithm = 'CSRT' # Default tracker
190+
if args.mil:
191+
tracker_algorithm = 'MIL'
192+
elif args.goturn:
193+
tracker_algorithm = 'GOTURN'
194+
elif args.dasiamrpn:
195+
tracker_algorithm = 'DaSiamRPN'
196+
elif args.nano:
197+
tracker_algorithm = 'Nano'
198+
elif args.vit:
199+
tracker_algorithm = 'Vit'
200+
elif args.csrt:
201+
tracker_algorithm = 'CSRT'
202+
elif args.kcf:
203+
tracker_algorithm = 'KCF'
204+
elif args.boosting:
205+
tracker_algorithm = 'Boosting'
206+
elif args.mosse:
207+
tracker_algorithm = 'MOSSE'
208+
elif args.medianflow:
209+
tracker_algorithm = 'MedianFlow'
210+
elif args.tld:
211+
tracker_algorithm = 'TLD'
212+
213+
logging.info(f"Selected Tracker Algorithm: {tracker_algorithm}")
214+
print("Selected Tracker Algorithm:", tracker_algorithm)
215+
216+
# Camera setup ###########################################################
217+
if isint(cap_device):
218+
cap_device = int(cap_device)
219+
cap = cv.VideoCapture(cap_device)
220+
if not cap.isOpened():
221+
logging.error(f"Cannot open video source: {cap_device}")
222+
sys.exit(1)
223+
cap.set(cv.CAP_PROP_FRAME_WIDTH, cap_width)
224+
cap.set(cv.CAP_PROP_FRAME_HEIGHT, cap_height)
225+
logging.info(f"Video capture started on {cap_device} with resolution {cap_width}x{cap_height}.")
226+
227+
# Load YOLOv8 model ######################################################
228+
try:
229+
model_path = r"D:\pycharm_projects\yolov8\runs\detect\drone_v9_300ep_32bath\weights\best.pt"
230+
model = YOLO(model_path, task='detect') # Ensure you have the correct path to your YOLOv8 model
231+
logging.info(f"YOLOv8 model loaded from {model_path}.")
232+
except Exception as e:
233+
logging.error(f"Failed to load YOLOv8 model: {e}")
234+
sys.exit(1)
235+
236+
# Queues for inter-thread communication ##################################
237+
frame_queue = queue.Queue(maxsize=5) # Increased maxsize to 5
238+
detection_queue = queue.Queue(maxsize=5)
239+
stop_event = threading.Event()
240+
241+
# Detection Thread ########################################################
242+
def detection_worker():
243+
while not stop_event.is_set():
244+
try:
245+
frame = frame_queue.get(timeout=1)
246+
logging.debug("Frame retrieved from frame_queue for detection.")
247+
detect_objects(model, frame, detection_queue)
248+
logging.debug("Object detection completed and results put into detection_queue.")
249+
except queue.Empty:
250+
continue
251+
except Exception as e:
252+
logging.error(f"Exception in detection_worker: {e}")
253+
254+
detection_thread = threading.Thread(target=detection_worker, daemon=True)
255+
detection_thread.start()
256+
logging.info("Detection thread started.")
257+
258+
# Tracker initialization #################################################
259+
window_name = 'Tracker Demo'
260+
cv.namedWindow(window_name)
261+
262+
tracker_info = None # Single TrackerInfo instance
263+
detected_bboxes = []
264+
frame_count = 0
265+
detection_interval = 5 # Perform detection every 5 frames
266+
267+
try:
268+
while cap.isOpened():
269+
ret, image = cap.read()
270+
if not ret:
271+
logging.info("No frame received. Exiting main loop.")
272+
break
273+
274+
debug_image = copy.deepcopy(image)
275+
frame_count += 1
276+
277+
# Put the frame into the frame_queue for detection every detection_interval frames
278+
if frame_count % detection_interval == 0:
279+
if frame_queue.full():
280+
try:
281+
removed_frame = frame_queue.get_nowait()
282+
logging.debug("Frame queue is full. Removed oldest frame to add a new one.")
283+
except queue.Empty:
284+
logging.warning("Frame queue was full but no frame to remove.")
285+
frame_queue.put(image)
286+
logging.debug("Frame added to frame_queue.")
287+
else:
288+
logging.debug(f"Skipping detection for frame {frame_count}.")
289+
290+
# Retrieve detection results if available
291+
try:
292+
detected_bboxes = detection_queue.get_nowait()
293+
if detected_bboxes and tracker_info is None:
294+
logging.debug(f"Retrieved {len(detected_bboxes)} bounding boxes from detection_queue.")
295+
# Initialize only the first detected bounding box
296+
initial_bbox = detected_bboxes[0]
297+
tracker_info = initialize_tracker(image, tracker_algorithm, initial_bbox)
298+
if tracker_info:
299+
# Assign a color
300+
tracker_info.color = color_list[0 % len(color_list)]
301+
logging.info(f"Initialized tracker '{tracker_algorithm}' with bbox: {initial_bbox}")
302+
except queue.Empty:
303+
pass # No new detections yet
304+
305+
if tracker_info:
306+
smoothed_bbox, ok = tracker_info.update(image)
307+
if ok and smoothed_bbox is not None:
308+
# Draw bounding box after smoothing
309+
cv.rectangle(debug_image,
310+
(smoothed_bbox[0], smoothed_bbox[1]),
311+
(smoothed_bbox[0] + smoothed_bbox[2], smoothed_bbox[1] + smoothed_bbox[3]),
312+
tracker_info.color,
313+
thickness=2)
314+
logging.debug(f"Tracker ({tracker_info.algorithm}) updated successfully with smoothed bbox: {smoothed_bbox}")
315+
else:
316+
# If tracking fails, remove the tracker_info
317+
logging.warning(f"Tracker ({tracker_info.algorithm}) failed to update. Removing tracker.")
318+
tracker_info = None
319+
320+
# Display processing time and inference speed
321+
# Replace static text with dynamic measurements if desired
322+
# For now, keeping it static as in the original code
323+
cv.putText(
324+
debug_image,
325+
f"Speed: 1.0ms preprocess, 17.9ms inference, 11.0ms postprocess per image at shape (1, 3, 512, 640)",
326+
(10, debug_image.shape[0] - 10),
327+
cv.FONT_HERSHEY_SIMPLEX,
328+
0.5,
329+
(0, 255, 255),
330+
1,
331+
cv.LINE_AA
332+
)
333+
334+
cv.imshow(window_name, debug_image)
335+
336+
k = cv.waitKey(1)
337+
if k == 32: # SPACE
338+
# Reinitialize tracker based on new selection
339+
try:
340+
new_detections = detection_queue.get_nowait()
341+
if new_detections:
342+
detected_bboxes = new_detections
343+
# Reinitialize tracker with the first detected bounding box
344+
initial_bbox = detected_bboxes[0]
345+
tracker_info = initialize_tracker(image, tracker_algorithm, initial_bbox)
346+
if tracker_info:
347+
tracker_info.color = color_list[0 % len(color_list)]
348+
logging.info(f"Re-initialized tracker '{tracker_algorithm}' with bbox: {initial_bbox}")
349+
except queue.Empty:
350+
logging.debug("No new detections available for reinitialization.")
351+
if k == 27: # ESC
352+
logging.info("ESC key pressed. Exiting.")
353+
break
354+
except KeyboardInterrupt:
355+
logging.info("KeyboardInterrupt received. Exiting.")
356+
except Exception as e:
357+
logging.error(f"Unexpected error in main loop: {e}")
358+
finally:
359+
# Cleanup ##############################################################
360+
stop_event.set()
361+
detection_thread.join(timeout=2)
362+
cap.release()
363+
cv.destroyAllWindows()
364+
logging.info("Resources released and program terminated gracefully.")
365+
366+
if __name__ == '__main__':
367+
main()

0 commit comments

Comments
(0)

AltStyle によって変換されたページ (->オリジナル) /