282 lines
8.8 KiB
Python
282 lines
8.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Created on Sun Mar 2 14:15:57 2025
|
||
|
||
@author: ym
|
||
"""
|
||
import numpy as np
|
||
import cv2
|
||
import os
|
||
from pathlib import Path
|
||
|
||
import sys
|
||
sys.path.append(r"D:\DetectTracking")
|
||
|
||
# from tracking.utils.read_data import extract_data_realtime, read_tracking_output_realtime
|
||
from tracking.utils.plotting import Annotator, colors
|
||
from tracking.utils import Boxes, IterableSimpleNamespace, yaml_load, boxes_add_fid
|
||
from tracking.trackers import BOTSORT, BYTETracker
|
||
from tracking.utils.showtrack import drawtracks
|
||
from hands.hand_inference import hand_pose
|
||
from tracking.utils.read_data import read_weight_sensor, extract_data_realtime, read_tracking_output_realtime
|
||
from contrast.feat_extract.config import config as conf
|
||
from contrast.feat_extract.inference import FeatsInterface
|
||
from tracking.utils.drawtracks import drawTrack
|
||
|
||
|
||
ReIDEncoder = FeatsInterface(conf)
|
||
|
||
|
||
W, H = 1024, 1280
|
||
Mode = 'front' #'back'
|
||
ImgFormat = ['.jpg', '.jpeg', '.png', '.bmp']
|
||
|
||
|
||
'''调用tracking()函数，利用本地跟踪算法获取各目标轨迹，可以比较本地跟踪算法与现场跟踪算法的区别。'''
|
||
def init_tracker(tracker_yaml = None, bs=1):
|
||
"""
|
||
Initialize tracker for object tracking during prediction.
|
||
"""
|
||
TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
|
||
cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))
|
||
|
||
tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
|
||
|
||
return tracker
|
||
|
||
|
||
def init_trackers(tracker_yaml = None, bs=1):
|
||
"""
|
||
Initialize trackers for object tracking during prediction.
|
||
"""
|
||
# tracker_yaml = r"./tracking/trackers/cfg/botsort.yaml"
|
||
|
||
TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
|
||
|
||
cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))
|
||
trackers = []
|
||
for _ in range(bs):
|
||
tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
|
||
trackers.append(tracker)
|
||
|
||
return trackers
|
||
|
||
|
||
def draw_box(img, tracks):
|
||
annotator = Annotator(img.copy(), line_width=2)
|
||
# for *xyxy, conf, cls in reversed(tracks):
|
||
# name = f'{int(cls)} {conf:.2f}'
|
||
# color = colors(int(cls), True)
|
||
# annotator.box_label(xyxy, name, color=color)
|
||
|
||
for *xyxy, id, conf, cls, fid, bid in reversed(tracks):
|
||
name = f'ID:{int(id)} {int(cls)} {conf:.2f}'
|
||
color = colors(int(cls), True)
|
||
annotator.box_label(xyxy, name, color=color)
|
||
|
||
|
||
|
||
im0 = annotator.result()
|
||
|
||
return im0
|
||
|
||
|
||
|
||
|
||
def tracking(bboxes, ffeats):
|
||
tracker_yaml = "./tracking/trackers/cfg/botsort.yaml"
|
||
|
||
tracker = init_tracker(tracker_yaml)
|
||
|
||
TrackBoxes = np.empty((0, 9), dtype = np.float32)
|
||
TracksDict = {}
|
||
|
||
frmIds = []
|
||
'''========================== 执行跟踪处理 ============================='''
|
||
# dets 与 feats 应保持严格对应
|
||
k=0
|
||
for dets, feats in zip(bboxes, ffeats):
|
||
|
||
frmIds.append(np.unique(dets[:, 6]).astype(np.int64)[0])
|
||
|
||
boxes = dets[:, :6]
|
||
det_tracking = Boxes(boxes).cpu().numpy()
|
||
tracks, outfeats = tracker.update(det_tracking, features=feats)
|
||
'''tracks: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]
|
||
0 1 2 3 4 5 6 7 8
|
||
这里，frame_index 也可以用视频的 帧ID 代替, box_index 保持不变
|
||
'''
|
||
k += 1
|
||
imgpath = r"D:\全实时\202502\tracker\Yolos_Tracking\tracker\1_1740891284792\1_1740891284792_{}.png".format(int(k))
|
||
img = cv2.imread(imgpath)
|
||
|
||
im0 = draw_box(img, tracks)
|
||
savepath = r"D:\全实时\202502\tracker\Yolos_Tracking\tracker\1_1740891284792\b\1_1740891284792_{}_b.png".format(k)
|
||
cv2.imwrite(savepath, im0)
|
||
|
||
|
||
if len(tracks):
|
||
TrackBoxes = np.concatenate([TrackBoxes, tracks], axis=0)
|
||
|
||
# =============================================================================
|
||
# FeatDict = {}
|
||
# for track in tracks:
|
||
# tid = int(track[8])
|
||
# FeatDict.update({tid: feats[tid, :]})
|
||
#
|
||
# frameID = tracks[0, 7]
|
||
#
|
||
# # print(f"frameID: {int(frameID)}")
|
||
# assert len(tracks) == len(FeatDict), f"Please check the func: tracker.update() at frameID({int(frameID)})"
|
||
#
|
||
# TracksDict[f"frame_{int(frameID)}"] = {"feats":FeatDict}
|
||
# =============================================================================
|
||
|
||
|
||
return TrackBoxes, TracksDict
|
||
|
||
def dotrack():
|
||
|
||
datapath = r"D:\全实时\202502\tracker\1_tracker_in.data"
|
||
|
||
bboxes, ffeats = extract_data_realtime(datapath)
|
||
trackerboxes, tracker_feat_dict = tracking(bboxes, ffeats)
|
||
|
||
print("done!")
|
||
|
||
|
||
# def plotbox():
|
||
|
||
# fpath = r"\\192.168.1.28\share\测试视频数据以及日志\全实时测试\V12\2025-3-3\20250303-103833-338_6928804010091_6928804010091\1_tracking_output.data"
|
||
# imgpath = r"D:\全实时\202502\result\Yolos_Tracking\20250303-103833-338_6928804010091_6928804010091\1_1740969517953"
|
||
# trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(fpath)
|
||
|
||
# for *xyxy, id, conf, cls, fid, bid in tracking_outboxes[0]:
|
||
# imgname = f"1_1740969517953_{int(fid)}.png"
|
||
|
||
# img_path = os.path.join(imgpath, imgname)
|
||
# img = cv2.imread(img_path)
|
||
# annotator = Annotator(img.copy(), line_width=2)
|
||
|
||
# name = f'ID:{int(id)} {int(cls)} {conf:.2f}'
|
||
# color = colors(int(cls), True)
|
||
# annotator.box_label(xyxy, name, color=color)
|
||
# im0 = annotator.result()
|
||
# cv2.imwrite(os.path.join(imgpath, f"1_1740969517953_{int(fid)}_.png"), im0)
|
||
|
||
# print(f"1_1740969676295_{int(fid)}_.png")
|
||
|
||
# print("done")
|
||
|
||
def video2imgs(videopath):
|
||
cap = cv2.VideoCapture(str(videopath))
|
||
k = 0
|
||
while True:
|
||
ret, frame = cap.read()
|
||
if frame is None:
|
||
break
|
||
k += 1
|
||
imgpath = videopath.parent / f"{videopath.stem}_{k}.png"
|
||
cv2.imwrite(str(imgpath), frame)
|
||
|
||
|
||
|
||
def extract_evtimgs(evtpath):
|
||
vidpaths = [v for v in evtpath.iterdir() if v.suffix == '.mp4']
|
||
for vidpath in vidpaths:
|
||
video2imgs(vidpath)
|
||
|
||
stamps = [name.stem.split('_')[1] for name in vidpaths]
|
||
|
||
if len(set(stamps)==1):
|
||
return stamps[0]
|
||
return None
|
||
|
||
def draw_tracking_boxes(evtpath, stamp):
|
||
for datapath in evtpath.iterdir():
|
||
if datapath.name.find('_tracking_output.data')<=0:
|
||
continue
|
||
|
||
camera = datapath.stem.split('_')[0]
|
||
trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(str(datapath))
|
||
|
||
## 该模块先读取轨迹数据，再根据帧ID读取相应图像
|
||
for *xyxy, id, conf, cls, fid, bid in tracking_outboxes[0]:
|
||
imgpath = evtpath / f"{camera}_{stamp}_{int(fid)}.png"
|
||
|
||
img = cv2.imread(str(imgpath))
|
||
annotator = Annotator(img.copy(), line_width=2)
|
||
|
||
name = f'ID:{int(id)} {int(cls)} {conf:.2f}'
|
||
color = colors(int(cls), True)
|
||
annotator.box_label(xyxy, name, color=color)
|
||
im0 = annotator.result()
|
||
cv2.imwrite(imgpath, im0)
|
||
|
||
print(datapath.name)
|
||
|
||
def draw_traj(evtpath):
|
||
for datapath in evtpath.iterdir():
|
||
if datapath.name.find('_tracking_output.data')<=0:
|
||
continue
|
||
|
||
fname = datapath.name
|
||
trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(datapath)
|
||
|
||
CamerType = fname.split('_')[0]
|
||
if CamerType == '1':
|
||
edgeline = cv2.imread("./CartTemp/board_ftmp_line.png")
|
||
if CamerType == '0':
|
||
edgeline = cv2.imread("./CartTemp/edgeline.png")
|
||
edgeline = drawTrack(tracking_outboxes, edgeline)
|
||
|
||
|
||
imgpath = datapath.parent / f"{datapath.stem}.png"
|
||
cv2.imwrite(str(imgpath), edgeline)
|
||
|
||
|
||
def main():
|
||
|
||
path = r"\\192.168.1.28\share\测试视频数据以及日志\全实时测试\V12\2025-3-3\20250303-104225-381_6920459958674"
|
||
evtpaths = [p for p in Path(path).iterdir() if p.is_dir()]
|
||
for evtpath in evtpaths:
|
||
#1. 从事件的前后摄视频提取图像
|
||
stamp = extract_evtimgs(evtpath)
|
||
|
||
#2. 根据 0/1_tracking_output.data 中提取的轨迹在img中绘制box
|
||
draw_tracking_boxes(evtpath, stamp)
|
||
|
||
#3. 根据 0/1_tracking_output.data 中提取的轨迹在edgeline中绘制box
|
||
draw_traj(evtpath)
|
||
|
||
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# dotrack()
|
||
# plotbox()
|
||
|
||
vpath = r"D:\datasets\ym\VID_20250307_105606"
|
||
extract_evtimgs(Path(vpath))
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|
||
|