202 lines
6.5 KiB
Python
202 lines
6.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Created on Sun Sep 29 08:59:21 2024
|
||
|
||
@author: ym
|
||
"""
|
||
import os
|
||
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
|
||
|
||
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]
|
||
|
||
input_enents.extend(input_enent)
|
||
|
||
|
||
return input_enents
|
||
|
||
def pipeline(
|
||
eventpath,
|
||
savepath,
|
||
SourceType = "image", # video
|
||
stdfeat_path = None
|
||
):
|
||
if SourceType == "video":
|
||
vpaths = get_video_pairs(eventpath)
|
||
elif SourceType == "image":
|
||
vpaths = get_image_pairs(eventpath)
|
||
|
||
'''======== 函数 yolo_resnet_tracker() 的参数字典 ========'''
|
||
opt = parse_opt()
|
||
optdict = vars(opt)
|
||
optdict["is_save_img"] = True
|
||
optdict["is_save_video"] = True
|
||
|
||
event_tracks = []
|
||
for vpath in vpaths:
|
||
'''事件结果文件夹'''
|
||
save_dir_event = Path(savepath) / Path(os.path.basename(eventpath))
|
||
if isinstance(vpath, list):
|
||
save_dir_video = save_dir_event / 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)
|
||
|
||
'''Yolo + Resnet + Tracker'''
|
||
optdict["source"] = vpath
|
||
optdict["save_dir"] = save_dir_video
|
||
|
||
|
||
tracksdict = yolo_resnet_tracker(**optdict)
|
||
|
||
bboxes = tracksdict['TrackBoxes']
|
||
|
||
bname = os.path.basename(vpath[0]) if isinstance(vpath, list) else os.path.basename(vpath)
|
||
if bname.split('_')[0] == "0" or bname.find('back')>=0:
|
||
vts = doBackTracks(bboxes, tracksdict)
|
||
vts.classify()
|
||
|
||
event_tracks.append(("back", vts))
|
||
|
||
if bname.split('_')[0] == "1" or bname.find('front')>=0:
|
||
vts = doFrontTracks(bboxes, tracksdict)
|
||
vts.classify()
|
||
event_tracks.append(("front", vts))
|
||
|
||
|
||
'''轨迹显示模块'''
|
||
illus = [None, None]
|
||
for CamerType, vts in event_tracks:
|
||
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, save_dir_event, CamerType, draw5p=True)
|
||
illus[0] = img_tracking
|
||
|
||
|
||
plt = plot_frameID_y2(vts)
|
||
plt.savefig(os.path.join(save_dir_event, "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, save_dir_event, 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(save_dir_event, "traj.png")
|
||
cv2.imwrite(trajpath, img_cat)
|
||
|
||
|
||
|
||
'''前后摄轨迹选择'''
|
||
|
||
if stdfeat_path is not None:
|
||
with open(stdfeat_path, 'rb') as f:
|
||
featDict = pickle.load(f)
|
||
|
||
|
||
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,
|
||
'''
|
||
parmDict = {}
|
||
parmDict["eventpath"] = r"\\192.168.1.28\share\测试视频数据以及日志\各模块测试记录\展厅测试\1120_展厅模型v801测试\扫A放A\20241121-144855-dce94b09-1100-43f1-92e8-33a1b538b159_6924743915848_6924743915848"
|
||
|
||
parmDict["savepath"] = r"D:\contrast\detect"
|
||
parmDict["SourceType"] = "image" # video, image
|
||
parmDict["stdfeat_path"] = None
|
||
|
||
pipeline(**parmDict)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|
||
|
||
|
||
|
||
|
||
|
||
|