lichen/detecttracking
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

modify pipeline.py

Browse Source
This commit is contained in:
王庆刚
2025-01-13 17:35:15 +08:00
parent 3d13b0d9c5
commit a50f777839
12 changed files with 322 additions and 211 deletions
Download Patch File Download Diff File Expand all files Collapse all files

209
pipeline.py
View File

@ -5,67 +5,48 @@ Created on Sun Sep 29 08:59:21 2024
@author: ym
"""
import os
import sys
# import sys
import cv2
import pickle
import argparse
import numpy as np
from pathlib import Path
from track_reid import parse_opt
from track_reid import yolo_resnet_tracker
# FILE = Path(__file__).resolve()
# ROOT = FILE.parents[0] # YOLOv5 root directory
# if str(ROOT) not in sys.path:
# sys.path.append(str(ROOT)) # add ROOT to PATH
# ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
from tracking.dotrack.dotracks_back import doBackTracks
from tracking.dotrack.dotracks_front import doFrontTracks
from tracking.utils.drawtracks import plot_frameID_y2, draw_all_trajectories
from utils.getsource import get_image_pairs, get_video_pairs
from tracking.utils.read_data import read_similar
def get_interbcd_inputenents():
bcdpath = r"\\192.168.1.28\share\测试_202406\contrast\std_barcodes_2192"
eventpath = r"\\192.168.1.28\share\测试_202406\0918"
barcodes = []
eventpaths = []
for featname in os.listdir(bcdpath):
barcode, ext = os.path.splitext(featname)
barcodes.append(barcode)
input_enents = []
for root, dirs, files in os.walk(eventpath):
input_enent = [os.path.join(root, d) for d in dirs if d.split('_')[-1] in barcodes]
def save_subimgs(imgdict, boxes, spath, ctype):
for i in range(len(boxes)):
fid, bid = int(boxes[i, 7]), int(boxes[i, 8])
if f"{fid}_{bid}" in imgdict.keys():
img = imgdict[f"{fid}_{bid}"]
imgpath = spath / f"{ctype}_{fid}_{bid}.png"
cv2.imwrite(imgpath, img)
input_enents.extend(input_enent)
return input_enents
def pipeline(
eventpath,
savepath = '',
SourceType = "image", # video
stdfeat_path = None
savepath,
SourceType,
weights
):
if SourceType == "video":
vpaths = get_video_pairs(eventpath)
elif SourceType == "image":
vpaths = get_image_pairs(eventpath)
'''
eventpath: 单个事件的存储路径
'''
'''======== 函数 yolo_resnet_tracker() 的参数字典 ========'''
opt = parse_opt()
optdict = vars(opt)
if SourceType == "video":
vpaths = get_video_pairs(eventpath)
elif SourceType == "image":
vpaths = get_image_pairs(eventpath)
optdict["weights"] = r'D:\DetectTracking\ckpts\best_cls10_0906.pt'
optdict["is_save_img"] = True
optdict["is_save_video"] = True
optdict = {}
optdict["weights"] = weights
event_tracks = []
@ -76,21 +57,35 @@ def pipeline(
and evtname.split('_')[-1].isdigit() else ''
'''事件结果存储文件夹'''
if not savepath:
savepath = Path(__file__).resolve().parents[0] / "evtresult"
savepath = Path(__file__).resolve().parents[0] / "events_result"
save_dir_event = Path(savepath) / evtname
pickpath = Path(savepath)/"pickfile"
if not pickpath.exists():
pickpath.mkdir(parents=True, exist_ok=True)
savepath_pipeline = Path(savepath) / Path("Yolos_Tracking") / evtname
"""ShoppingDict pickle 文件保存地址 """
savepath_spdict = Path(savepath) / "ShoppingDict_pkfile"
if not savepath_spdict.exists():
savepath_spdict.mkdir(parents=True, exist_ok=True)
pf_path = Path(savepath_spdict) / Path(str(evtname)+".pickle")
# if pf_path.exists():
# return
ShoppingDict = {"eventPath": eventpath,
"eventName": evtname,
"barcode": barcode,
"eventType": '', # "input", "output", "other"
"frontCamera": {},
"backCamera": {}}
"backCamera": {},
"one2n": []
}
procpath = Path(eventpath).joinpath('process.data')
if procpath.is_file():
SimiDict = read_similar(procpath)
ShoppingDict["one2n"] = SimiDict['one2n']
for vpath in vpaths:
'''相机事件字典构造'''
@ -115,16 +110,20 @@ def pipeline(
'''事件结果存储文件夹'''
if isinstance(vpath, list):
save_dir_video = save_dir_event / Path("images")
savepath_pipeline_imgs = savepath_pipeline / Path("images")
else:
save_dir_video = save_dir_event / Path(str(Path(vpath).stem))
if not save_dir_video.exists():
save_dir_video.mkdir(parents=True, exist_ok=True)
savepath_pipeline_imgs = savepath_pipeline / Path(str(Path(vpath).stem))
if not savepath_pipeline_imgs.exists():
savepath_pipeline_imgs.mkdir(parents=True, exist_ok=True)
savepath_pipeline_subimgs = savepath_pipeline / Path("subimgs")
if not savepath_pipeline_subimgs.exists():
savepath_pipeline_subimgs.mkdir(parents=True, exist_ok=True)
'''Yolo + Resnet + Tracker'''
optdict["source"] = vpath
optdict["save_dir"] = save_dir_video
optdict["save_dir"] = savepath_pipeline_imgs
yrtOut = yolo_resnet_tracker(**optdict)
@ -164,13 +163,30 @@ def pipeline(
CameraEvent["tracking"] = vts
ShoppingDict["frontCamera"] = CameraEvent
# pklpath = save_dir_event / "ShoppingDict.pkl"
# with open(str(pklpath), 'wb') as f:
# pickle.dump(ShoppingDict, f)
pf_path = Path(pickpath) / Path(str(evtname)+".pkl")
with open(str(pf_path), 'wb') as f:
pickle.dump(ShoppingDict, f)
for CamerType, vts in event_tracks:
if len(vts.tracks)==0: continue
if CamerType == 'front':
yolos = ShoppingDict["frontCamera"]["yoloResnetTracker"]
ctype = 1
if CamerType == 'back':
yolos = ShoppingDict["backCamera"]["yoloResnetTracker"]
ctype = 0
imgdict = {}
for y in yolos:
imgdict.update(y["imgs"])
for track in vts.Residual:
if isinstance(track, np.ndarray):
save_subimgs(imgdict, track, savepath_pipeline_subimgs, ctype)
else:
save_subimgs(imgdict, track.boxes, savepath_pipeline_subimgs, ctype)
'''轨迹显示模块'''
illus = [None, None]
@ -181,24 +197,23 @@ def pipeline(
edgeline = cv2.imread("./tracking/shopcart/cart_tempt/board_ftmp_line.png")
h, w = edgeline.shape[:2]
nh, nw = h//2, w//2
edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA)
# nh, nw = h//2, w//2
# edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA)
img_tracking = draw_all_trajectories(vts, edgeline, save_dir_event, CamerType, draw5p=True)
img_tracking = draw_all_trajectories(vts, edgeline, savepath_pipeline, CamerType, draw5p=True)
illus[0] = img_tracking
plt = plot_frameID_y2(vts)
plt.savefig(os.path.join(save_dir_event, "front_y2.png"))
plt.savefig(os.path.join(savepath_pipeline, "front_y2.png"))
if CamerType == 'back':
edgeline = cv2.imread("./tracking/shopcart/cart_tempt/edgeline.png")
h, w = edgeline.shape[:2]
nh, nw = h//2, w//2
edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA)
# nh, nw = h//2, w//2
# edgeline = cv2.resize(edgeline, (nw, nh), interpolation=cv2.INTER_AREA)
img_tracking = draw_all_trajectories(vts, edgeline, save_dir_event, CamerType, draw5p=True)
img_tracking = draw_all_trajectories(vts, edgeline, savepath_pipeline, CamerType, draw5p=True)
illus[1] = img_tracking
illus = [im for im in illus if im is not None]
@ -208,64 +223,25 @@ def pipeline(
H, W = img_cat.shape[:2]
cv2.line(img_cat, (int(W/2), 0), (int(W/2), int(H)), (128, 128, 255), 3)
trajpath = os.path.join(save_dir_event, "traj.png")
trajpath = os.path.join(savepath_pipeline, "trajectory.png")
cv2.imwrite(trajpath, img_cat)
def main_loop():
bcdpath = r"\\192.168.1.28\share\测试_202406\contrast\std_barcodes_2192"
eventpath = r"\\192.168.1.28\share\测试_202406\0918\images1"
SourceType = "image" # video, image
barcodes = []
input_enents = []
output_events = []
'''1. 获得barcode标准特征集列表'''
for featname in os.listdir(bcdpath):
barcode, ext = os.path.splitext(featname)
if not barcode.isdigit() or len(barcode)<=8 or ext != ".pickle" :
continue
barcodes.append(barcode)
'''2. 构造（放入事件，标准特征）对'''
for filename in os.listdir(eventpath):
'''barcode为时间文件夹的最后一个字段'''
bcd = filename.split('_')[-1]
event_path = os.path.join(eventpath, filename)
stdfeat_path = None
if bcd in barcodes:
stdfeat_path = os.path.join(bcdpath, f"{bcd}.pickle")
input_enents.append((event_path, stdfeat_path))
parmDict = {}
parmDict["SourceType"] = "image"
parmDict["savepath"] = r"D:\contrast\detect"
for eventpath, stdfeat_path in input_enents:
parmDict["eventpath"] = eventpath
parmDict["stdfeat_path"] = stdfeat_path
pipeline(**parmDict)
def main():
'''
函数：pipeline()，遍历事件文件夹，选择类型 image 或 video,
'''
evtdir = r"\\192.168.1.28\share\测试视频数据以及日志\算法全流程测试\202412\images"
evtdir = Path(evtdir)
parmDict = {}
parmDict["savepath"] = r"D:\contrast\202412测试"
evtdir = r"\\192.168.1.28\share\测试视频数据以及日志\算法全流程测试\202412\images"
parmDict["SourceType"] = "video" # video, image
parmDict["stdfeat_path"] = None
k = 0
errEvents = []
parmDict["savepath"] = r"\\192.168.1.28\share\测试视频数据以及日志\算法全流程测试\202412\result"
parmDict["weights"] = r'D:\DetectTracking\ckpts\best_cls10_0906.pt'
evtdir = Path(evtdir)
k, errEvents = 0, []
for item in evtdir.iterdir():
if item.is_dir():
# item = r"D:\exhibition\images\images2\images2"
if item.is_dir():
# item = evtdir/Path("20241209-160201-b97f7a0e-7322-4375-9f17-c475500097e9_6926265317292")
parmDict["eventpath"] = item
# pipeline(**parmDict)
@ -273,10 +249,9 @@ def main():
pipeline(**parmDict)
except Exception as e:
errEvents.append(str(item))
# k+=1
# if k==1:
# break
k+=1
if k==1:
break
errfile = os.path.join(parmDict["savepath"], f'error_events.txt')
with open(errfile, 'w', encoding='utf-8') as f: