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 6572e3f

Browse files
committed
tracker work paralell with detector, work var.
chsnge req in ultralytics lib ( because need for tracker spec ver of opencv )
1 parent fda4a40 commit 6572e3f

File tree

1 file changed

+360
-0
lines changed

1 file changed

+360
-0
lines changed

‎work_paralel.py

Lines changed: 360 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,360 @@
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+
12+
import cv2 as cv
13+
import numpy as np
14+
from ultralytics import YOLO
15+
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+
28+
def get_args():
29+
parser = argparse.ArgumentParser()
30+
31+
parser.add_argument("--device", default="sample_movie/bird.mp4")
32+
parser.add_argument("--width", help='cap width', type=int, default=1280)
33+
parser.add_argument("--height", help='cap height', type=int, default=720)
34+
35+
parser.add_argument('--use_mil', action='store_true')
36+
parser.add_argument('--use_goturn', action='store_true')
37+
parser.add_argument('--use_dasiamrpn', action='store_true')
38+
parser.add_argument('--use_csrt', action='store_true')
39+
parser.add_argument('--use_kcf', action='store_true')
40+
parser.add_argument('--use_boosting', action='store_true')
41+
parser.add_argument('--use_mosse', action='store_true')
42+
parser.add_argument('--use_medianflow', action='store_true')
43+
parser.add_argument('--use_tld', action='store_true')
44+
parser.add_argument('--use_nano', action='store_true')
45+
parser.add_argument('--use_vit', action='store_true')
46+
47+
args = parser.parse_args()
48+
logging.debug(f"Parsed arg: {args}")
49+
50+
return args
51+
52+
53+
def isint(s):
54+
p = '[-+]?\d+'
55+
logging.debug(f"Checking {p}")
56+
return True if re.fullmatch(p, s) else False
57+
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+
76+
def initialize_tracker_list(image, tracker_algorithm_list, detected_bboxes):
77+
tracker_list = []
78+
79+
# Tracker list generation
80+
for tracker_algorithm in tracker_algorithm_list:
81+
for bbox in detected_bboxes:
82+
tracker = None
83+
logging.debug(f"Initializing tracker '{tracker_algorithm}' with bbox: {bbox}")
84+
if tracker_algorithm == 'MIL':
85+
tracker = cv.TrackerMIL_create()
86+
elif tracker_algorithm == 'GOTURN':
87+
params = cv.TrackerGOTURN_Params()
88+
params.modelTxt = "model/GOTURN/goturn.prototxt"
89+
params.modelBin = "model/GOTURN/goturn.caffemodel"
90+
tracker = cv.TrackerGOTURN_create(params)
91+
elif tracker_algorithm == 'DaSiamRPN':
92+
params = cv.TrackerDaSiamRPN_Params()
93+
params.model = "model/DaSiamRPN/dasiamrpn_model.onnx"
94+
params.kernel_r1 = "model/DaSiamRPN/dasiamrpn_kernel_r1.onnx"
95+
params.kernel_cls1 = "model/DaSiamRPN/dasiamrpn_kernel_cls1.onnx"
96+
tracker = cv.TrackerDaSiamRPN_create(params)
97+
elif tracker_algorithm == 'Nano':
98+
params = cv.TrackerNano_Params()
99+
params.backbone = "model/nanotrackv2/nanotrack_backbone_sim.onnx"
100+
params.neckhead = "model/nanotrackv2/nanotrack_head_sim.onnx"
101+
tracker = cv.TrackerNano_create(params)
102+
elif tracker_algorithm == 'Vit':
103+
params = cv.TrackerVit_Params()
104+
params.net = "model/vit/object_tracking_vittrack_2023sep.onnx"
105+
tracker = cv.TrackerVit_create(params)
106+
elif tracker_algorithm == 'CSRT':
107+
tracker = cv.TrackerCSRT_create()
108+
elif tracker_algorithm == 'KCF':
109+
tracker = cv.TrackerKCF_create()
110+
elif tracker_algorithm == 'Boosting':
111+
tracker = cv.legacy_TrackerBoosting.create()
112+
elif tracker_algorithm == 'MOSSE':
113+
tracker = cv.legacy_TrackerMOSSE.create()
114+
elif tracker_algorithm == 'MedianFlow':
115+
tracker = cv.legacy_TrackerMedianFlow.create()
116+
elif tracker_algorithm == 'TLD':
117+
tracker = cv.legacy_TrackerTLD.create()
118+
else:
119+
logging.warning(f"Unknown tracker algorithm: {tracker_algorithm}")
120+
121+
if tracker is not None:
122+
try:
123+
tracker.init(image, bbox)
124+
tracker_list.append(tracker)
125+
logging.debug(f"Successfully initialized '{tracker_algorithm}' tracker with bbox: {bbox}")
126+
except Exception as e:
127+
logging.error(f"Exception during tracker initialization for '{tracker_algorithm}' with bbox {bbox}: {e}")
128+
else:
129+
logging.error(f"Failed to initialize '{tracker_algorithm}' tracker with bbox: {bbox}")
130+
131+
logging.info(f"Total trackers initialized: {len(tracker_list)}")
132+
return tracker_list
133+
134+
135+
def main():
136+
color_list = [
137+
[255, 0, 0], # blue
138+
[255, 255, 0], # aqua
139+
[0, 255, 0], # lime
140+
[128, 0, 128], # purple
141+
[0, 0, 255], # red
142+
[255, 0, 255], # fuchsia
143+
[0, 128, 0], # green
144+
[128, 128, 0], # teal
145+
[0, 0, 128], # maroon
146+
[0, 128, 128], # olive
147+
[0, 255, 255], # yellow
148+
]
149+
150+
# Parse arguments ########################################################
151+
args = get_args()
152+
153+
cap_device = args.device
154+
cap_width = args.width
155+
cap_height = args.height
156+
157+
use_mil = args.use_mil
158+
use_goturn = args.use_goturn
159+
use_dasiamrpn = args.use_dasiamrpn
160+
use_csrt = args.use_csrt
161+
use_kcf = args.use_kcf
162+
use_boosting = args.use_boosting
163+
use_mosse = args.use_mosse
164+
use_medianflow = args.use_medianflow
165+
use_tld = args.use_tld
166+
use_nano = args.use_nano
167+
use_vit = args.use_vit
168+
169+
# Tracker algorithm selection ############################################
170+
tracker_algorithm_list = []
171+
if use_mil:
172+
tracker_algorithm_list.append('MIL')
173+
if use_goturn:
174+
tracker_algorithm_list.append('GOTURN')
175+
if use_dasiamrpn:
176+
tracker_algorithm_list.append('DaSiamRPN')
177+
if use_csrt:
178+
tracker_algorithm_list.append('CSRT')
179+
if use_kcf:
180+
tracker_algorithm_list.append('KCF')
181+
if use_boosting:
182+
tracker_algorithm_list.append('Boosting')
183+
if use_mosse:
184+
tracker_algorithm_list.append('MOSSE')
185+
if use_medianflow:
186+
tracker_algorithm_list.append('MedianFlow')
187+
if use_tld:
188+
tracker_algorithm_list.append('TLD')
189+
if use_nano:
190+
tracker_algorithm_list.append('Nano')
191+
if use_vit:
192+
tracker_algorithm_list.append('Vit')
193+
194+
if len(tracker_algorithm_list) == 0:
195+
tracker_algorithm_list.append('DaSiamRPN')
196+
logging.info(f"Selected Tracker Algorithms: {tracker_algorithm_list}")
197+
print("Selected Tracker Algorithms:", tracker_algorithm_list)
198+
199+
# Camera setup ###########################################################
200+
if isint(cap_device):
201+
cap_device = int(cap_device)
202+
cap = cv.VideoCapture(cap_device)
203+
if not cap.isOpened():
204+
logging.error(f"Cannot open video source: {cap_device}")
205+
sys.exit(1)
206+
cap.set(cv.CAP_PROP_FRAME_WIDTH, cap_width)
207+
cap.set(cv.CAP_PROP_FRAME_HEIGHT, cap_height)
208+
logging.info(f"Video capture started on {cap_device} with resolution {cap_width}x{cap_height}.")
209+
210+
# Load YOLOv8 model ######################################################
211+
try:
212+
model_path = r"D:\pycharm_projects\yolov8\runs\detect\drone_v9_300ep_32bath\weights\best.pt"
213+
model = YOLO(model_path, task='detect') # Ensure you have the correct path to your YOLOv8 model
214+
logging.info(f"YOLOv8 model loaded from {model_path}.")
215+
except Exception as e:
216+
logging.error(f"Failed to load YOLOv8 model: {e}")
217+
sys.exit(1)
218+
219+
# Queues for inter-thread communication ##################################
220+
frame_queue = queue.Queue(maxsize=1)
221+
detection_queue = queue.Queue(maxsize=1)
222+
stop_event = threading.Event()
223+
224+
# Detection Thread ########################################################
225+
def detection_worker():
226+
while not stop_event.is_set():
227+
try:
228+
frame = frame_queue.get(timeout=1)
229+
logging.debug("Frame retrieved from frame_queue for detection.")
230+
detect_objects(model, frame, detection_queue)
231+
logging.debug("Object detection completed and results put into detection_queue.")
232+
except queue.Empty:
233+
continue
234+
except Exception as e:
235+
logging.error(f"Exception in detected_worker: {e}")
236+
237+
detection_thread = threading.Thread(target=detection_worker, daemon=True)
238+
detection_thread.start()
239+
logging.info("Detection thread started.")
240+
241+
# Tracker initialization #################################################
242+
window_name = 'Tracker Demo'
243+
cv.namedWindow(window_name)
244+
245+
tracker_list = []
246+
detected_bboxes = []
247+
248+
try:
249+
while cap.isOpened():
250+
ret, image = cap.read()
251+
if not ret:
252+
logging.info("No frame received. Exiting main loop.")
253+
break
254+
255+
debug_image = copy.deepcopy(image)
256+
257+
# Put the frame into the frame_queue for detection
258+
if not frame_queue.full():
259+
frame_queue.put(image)
260+
logging.debug("Frame added to frame_queue.")
261+
else:
262+
logging.debug("Frame queue is full. Skipping frame.")
263+
264+
# Retrieve detection results if available
265+
try:
266+
detected_bboxes = detection_queue.get_nowait()
267+
if detected_bboxes:
268+
logging.debug(f"Retrieved {len(detected_bboxes)} bounding boxes from detection_queue.")
269+
tracker_list = initialize_tracker_list(image, tracker_algorithm_list, detected_bboxes)
270+
logging.info(f"Initialized {len(tracker_list)} trackers based on detected bounding boxes.")
271+
except queue.Empty:
272+
pass # No new detections yet
273+
274+
elapsed_time_list = []
275+
tracker_scores = [] # Initialize a list to store tracker scores
276+
277+
for index, tracker in enumerate(tracker_list):
278+
# Update tracking
279+
start_time = time.time()
280+
ok, bbox = tracker.update(image)
281+
try:
282+
tracker_score = tracker.getTrackingScore()
283+
except:
284+
tracker_score = '-'
285+
286+
elapsed_time = time.time() - start_time
287+
elapsed_time_list.append(elapsed_time)
288+
tracker_scores.append(tracker_score) # Append the score to the list
289+
290+
if ok:
291+
# Draw bounding box after tracking
292+
new_bbox = [
293+
int(bbox[0]),
294+
int(bbox[1]),
295+
int(bbox[2]),
296+
int(bbox[3])
297+
]
298+
cv.rectangle(debug_image,
299+
(new_bbox[0], new_bbox[1]),
300+
(new_bbox[0] + new_bbox[2], new_bbox[1] + new_bbox[3]),
301+
color_list[index % len(color_list)],
302+
thickness=2)
303+
logging.debug(f"Tracker {index} updated successfully with bbox: {new_bbox}")
304+
else:
305+
# If tracking fails, reset trackers
306+
logging.warning(f"Tracker {index} failed to update. Resetting all trackers.")
307+
tracker_list = []
308+
break
309+
310+
# Display processing time and tracker scores for each algorithm
311+
for index, tracker_algorithm in enumerate(tracker_algorithm_list):
312+
if index < len(elapsed_time_list):
313+
elapsed_time_ms = elapsed_time_list[index] * 1000
314+
if index < len(tracker_scores):
315+
score = tracker_scores[index]
316+
if score != '-':
317+
text = f"{tracker_algorithm} : {elapsed_time_ms:.1f}ms Score:{score:.2f}"
318+
else:
319+
text = f"{tracker_algorithm} : {elapsed_time_ms:.1f}ms"
320+
else:
321+
text = f"{tracker_algorithm} : {elapsed_time_ms:.1f}ms"
322+
else:
323+
text = f"{tracker_algorithm} : N/A"
324+
325+
cv.putText(
326+
debug_image,
327+
text,
328+
(10, int(25 * (index + 1))),
329+
cv.FONT_HERSHEY_SIMPLEX,
330+
0.7,
331+
color_list[index % len(color_list)],
332+
2,
333+
cv.LINE_AA
334+
)
335+
336+
cv.imshow(window_name, debug_image)
337+
338+
k = cv.waitKey(1)
339+
if k == 32: # SPACE
340+
# Reinitialize trackers based on new selection
341+
detected_bboxes = detection_queue.get()
342+
tracker_list = initialize_tracker_list(image, tracker_algorithm_list, detected_bboxes)
343+
if k == 27: # ESC
344+
logging.info("ESC key pressed. Exiting.")
345+
break
346+
except KeyboardInterrupt:
347+
logging.info("KeyboardInterrupt received. Exiting.")
348+
except Exception as e:
349+
logging.error(f"Unexpected error in main loop: {e}")
350+
finally:
351+
# Cleanup ##############################################################
352+
stop_event.set()
353+
detection_thread.join(timeout=2)
354+
cap.release()
355+
cv.destroyAllWindows()
356+
logging.info("Resources released and program terminated gracefully.")
357+
358+
359+
if __name__ == '__main__':
360+
main()

0 commit comments

Comments
(0)

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