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

增加了单帧入侵判断及yoloV10

Browse Source
This commit is contained in:
18262620154
2025-04-11 17:02:39 +08:00
parent 798c596acc
commit e044c85a04
197 changed files with 1863 additions and 997 deletions
Download Patch File Download Diff File Expand all files Collapse all files

313
pipeline.py
View File

@ -60,48 +60,158 @@ def save_subimgs_1(imgdict, boxes, spath, ctype, simidict = None):
cv2.imwrite(imgpath, img)
def show_result(event_tracks, yrtDict, savepath_pipe):
'''保存 Tracking 输出的运动轨迹子图，并记录相似度'''
savepath_pipe_subimgs = savepath_pipe / Path("subimgs")
if not savepath_pipe_subimgs.exists():
savepath_pipe_subimgs.mkdir(parents=True, exist_ok=True)
for CamerType, vts in event_tracks:
if len(vts.tracks)==0: continue
if CamerType == 'front':
# yolos = ShoppingDict["frontCamera"]["yoloResnetTracker"]
yolos = yrtDict["frontyrt"]
ctype = 1
if CamerType == 'back':
# yolos = ShoppingDict["backCamera"]["yoloResnetTracker"]
yolos = yrtDict["backyrt"]
ctype = 0
imgdict, featdict, simidict = {}, {}, {}
for y in yolos:
imgdict.update(y["imgs"])
featdict.update(y["feats"])
simidict.update(y["featsimi"])
def pipeline(
eventpath,
savepath,
SourceType,
weights,
YoloVersion="V5"
):
'''
eventpath: 单个事件的存储路径
for track in vts.Residual:
if isinstance(track, np.ndarray):
save_subimgs(imgdict, track, savepath_pipe_subimgs, ctype, featdict)
else:
save_subimgs(imgdict, track.slt_boxes, savepath_pipe_subimgs, ctype, featdict)
'''(3) 轨迹显示与保存'''
illus = [None, None]
for CamerType, vts in event_tracks:
if len(vts.tracks)==0: continue
if CamerType == 'front':
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)
img_tracking = draw_all_trajectories(vts, edgeline, savepath_pipe, CamerType, draw5p=True)
illus[0] = img_tracking
plt = plot_frameID_y2(vts)
plt.savefig(os.path.join(savepath_pipe, "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)
img_tracking = draw_all_trajectories(vts, edgeline, savepath_pipe, CamerType, draw5p=True)
illus[1] = img_tracking
'''
optdict = {}
optdict["weights"] = weights
if SourceType == "video":
vpaths = get_video_pairs(eventpath)
elif SourceType == "image":
vpaths = get_image_pairs(eventpath)
event_tracks = []
illus = [im for im in illus if im is not None]
if len(illus):
img_cat = np.concatenate(illus, axis = 1)
if len(illus)==2:
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(savepath_pipe, "trajectory.png")
cv2.imwrite(trajpath, img_cat)
def pipeline(eventpath,
SourceType,
weights,
DataType = "raw", #raw, pkl: images or videos, pkl, pickle file
YoloVersion="V5",
savepath = None,
saveimages = True
):
## 构造购物事件字典
evtname = Path(eventpath).stem
barcode = evtname.split('_')[-1] if len(evtname.split('_'))>=2 \
and len(evtname.split('_')[-1])>=8 \
and evtname.split('_')[-1].isdigit() else ''
'''事件结果存储文件夹'''
'''事件结果存储文件夹: savepath_pipe, savepath_pkl'''
if not savepath:
savepath = Path(__file__).resolve().parents[0] / "events_result"
savepath_pipe = Path(savepath) / Path("yolos_tracking") / evtname
savepath_pipeline = Path(savepath) / Path("Yolos_Tracking") / evtname
savepath_pkl = Path(savepath) / "shopping_pkl"
if not savepath_pkl.exists():
savepath_pkl.mkdir(parents=True, exist_ok=True)
pklpath = Path(savepath_pkl) / Path(str(evtname)+".pickle")
"""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")
yrt_out = []
if DataType == "raw":
### 不重复执行已经过yolo-resnet-tracker
if pklpath.exists():
print(f"Pickle file have saved: {evtname}.pickle")
return
# if pf_path.exists():
# print(f"Pickle file have saved: {evtname}.pickle")
# return
if SourceType == "video":
vpaths = get_video_pairs(eventpath)
elif SourceType == "image":
vpaths = get_image_pairs(eventpath)
for vpath in vpaths:
'''================= 2. 事件结果存储文件夹 ================='''
if isinstance(vpath, list):
savepath_pipe_imgs = savepath_pipe / Path("images")
else:
savepath_pipe_imgs = savepath_pipe / Path(str(Path(vpath).stem))
if not savepath_pipe_imgs.exists():
savepath_pipe_imgs.mkdir(parents=True, exist_ok=True)
optdict = {}
optdict["weights"] = weights
optdict["source"] = vpath
optdict["save_dir"] = savepath_pipe_imgs
optdict["is_save_img"] = saveimages
optdict["is_save_video"] = True
if YoloVersion == "V5":
yrtOut = yolo_resnet_tracker(**optdict)
elif YoloVersion == "V10":
yrtOut = yolov10_resnet_tracker(**optdict)
yrt_out.append((vpath, yrtOut))
elif DataType == "pkl":
pass
else:
return
'''====================== 构造 ShoppingDict 模块 ======================='''
ShoppingDict = {"eventPath": eventpath,
@ -112,16 +222,14 @@ def pipeline(
"backCamera": {},
"one2n": [] #
}
yrtDict = {}
procpath = Path(eventpath).joinpath('process.data')
if procpath.is_file():
SimiDict = read_similar(procpath)
ShoppingDict["one2n"] = SimiDict['one2n']
for vpath in vpaths:
yrtDict = {}
event_tracks = []
for vpath, yrtOut in yrt_out:
'''================= 1. 构造相机事件字典 ================='''
CameraEvent = {"cameraType": '', # "front", "back"
"videoPath": '',
@ -140,34 +248,10 @@ def pipeline(
CameraEvent["cameraType"] = "back"
if bname.split('_')[0] == "1" or bname.find('front')>=0:
CameraEvent["cameraType"] = "front"
'''================= 2. 事件结果存储文件夹 ================='''
if isinstance(vpath, list):
savepath_pipeline_imgs = savepath_pipeline / Path("images")
else:
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)
'''================= 3. Yolo + Resnet + Tracker ================='''
optdict["source"] = vpath
optdict["save_dir"] = savepath_pipeline_imgs
optdict["is_save_img"] = True
optdict["is_save_video"] = True
if YoloVersion == "V5":
yrtOut = yolo_resnet_tracker(**optdict)
elif YoloVersion == "V10":
yrtOut = yolov10_resnet_tracker(**optdict)
'''2种保存方式: (1) no save subimg, (2) save img'''
###(1) save images
yrtOut_save = []
for frdict in yrtOut:
fr_dict = {}
@ -177,6 +261,7 @@ def pipeline(
yrtOut_save.append(fr_dict)
CameraEvent["yoloResnetTracker"] = yrtOut_save
###(2) no save images
# CameraEvent["yoloResnetTracker"] = yrtOut
'''================= 4. tracking ================='''
@ -219,108 +304,58 @@ def pipeline(
yrtDict["frontyrt"] = yrtOut
'''========================== 保存模块 ================================='''
'''(1) 保存 ShoppingDict 事件'''
with open(str(pf_path), 'wb') as f:
# 保存 ShoppingDict
with open(str(pklpath), 'wb') as f:
pickle.dump(ShoppingDict, f)
'''(2) 保存 Tracking 输出的运动轨迹子图，并记录相似度'''
for CamerType, vts in event_tracks:
if len(vts.tracks)==0: continue
if CamerType == 'front':
# yolos = ShoppingDict["frontCamera"]["yoloResnetTracker"]
yolos = yrtDict["frontyrt"]
ctype = 1
if CamerType == 'back':
# yolos = ShoppingDict["backCamera"]["yoloResnetTracker"]
yolos = yrtDict["backyrt"]
ctype = 0
imgdict, featdict, simidict = {}, {}, {}
for y in yolos:
imgdict.update(y["imgs"])
featdict.update(y["feats"])
simidict.update(y["featsimi"])
for track in vts.Residual:
if isinstance(track, np.ndarray):
save_subimgs(imgdict, track, savepath_pipeline_subimgs, ctype, featdict)
else:
save_subimgs(imgdict, track.slt_boxes, savepath_pipeline_subimgs, ctype, featdict)
'''(3) 轨迹显示与保存'''
illus = [None, None]
for CamerType, vts in event_tracks:
if len(vts.tracks)==0: continue
if CamerType == 'front':
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)
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(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)
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]
if len(illus):
img_cat = np.concatenate(illus, axis = 1)
if len(illus)==2:
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(savepath_pipeline, "trajectory.png")
cv2.imwrite(trajpath, img_cat)
# 绘制并保存轨迹图
show_result(event_tracks, yrtDict, savepath_pipe)
def execute_pipeline(evtdir = r"D:\datasets\ym\后台数据\unzip",
source_type = "video", # video, image,
DataType = "raw", # raw, pkl
save_path = r"D:\work\result_pipeline",
kk=1,
source_type = "video", # video, image,
yolo_ver = "V10", # V10, V5
weight_yolo_v5 = r'./ckpts/best_cls10_0906.pt' ,
weight_yolo_v10 = r'./ckpts/best_v10s_width0375_1205.pt',
k=0
saveimages = True
):
'''
运行函数 pipeline()，遍历事件文件夹，每个文件夹是一个事件
'''
parmDict = {}
parmDict["SourceType"] = source_type
parmDict["DataType"] = DataType
parmDict["savepath"] = save_path
parmDict["SourceType"] = source_type
parmDict["YoloVersion"] = yolo_ver
if parmDict["YoloVersion"] == "V5":
parmDict["weights"] = weight_yolo_v5
elif parmDict["YoloVersion"] == "V10":
parmDict["weights"] = weight_yolo_v10
parmDict["saveimages"] = saveimages
evtdir = Path(evtdir)
errEvents = []
k = 0
for item in evtdir.iterdir():
if item.is_dir():
item = evtdir/Path("20250310-175352-741")
# item = evtdir/Path("20241212-171505-f0afe929-fdfe-4efa-94d0-2fa748d65fbb_6907992518930")
parmDict["eventpath"] = item
pipeline(**parmDict)
# try:
# pipeline(**parmDict)
# except Exception as e:
# errEvents.append(str(item))
k+=1
if k==1:
if kk is not None and k==kk:
break
errfile = os.path.join(parmDict["savepath"], 'error_events.txt')
@ -329,12 +364,20 @@ def execute_pipeline(evtdir = r"D:\datasets\ym\后台数据\unzip",
f.write(line + '\n')
if __name__ == "__main__":
execute_pipeline()
datapath = r'/home/wqg/dataset/test_dataset/base_dataset/single_event/source/'
savepath = r'/home/wqg/dataset/pipeline/test_result/single_event_V10'
execute_pipeline(evtdir = datapath,
DataType = "raw", # raw, pkl
kk=1,
source_type = "video", # video, image,
save_path = savepath,
yolo_ver = "V10", # V10, V5
weight_yolo_v5 = r'./ckpts/best_cls10_0906.pt' ,
weight_yolo_v10 = r'./ckpts/best_v10s_width0375_1205.pt',
saveimages = False
)
# spath_v10 = r"D:\work\result_pipeline_v10"
# spath_v5 = r"D:\work\result_pipeline_v5"
# execute_pipeline(save_path=spath_v10, yolo_ver="V10")
# execute_pipeline(save_path=spath_v5, yolo_ver="V5")