# -*- coding: utf-8 -*- """ Created on Tue Nov 26 17:35:05 2024 @author: ym """ import os import cv2 import pickle import numpy as np from pathlib import Path import sys sys.path.append(r"D:\DetectTracking") from tracking.utils.plotting import Annotator, colors from tracking.utils.drawtracks import drawTrack from tracking.utils.read_data import extract_data, read_tracking_output, read_similar from tracking.utils.read_data import extract_data_realtime, read_tracking_output_realtime IMG_FORMAT = ['.bmp', '.jpg', '.jpeg', '.png'] VID_FORMAT = ['.mp4', '.avi'] def save_data(event, resultPath=None): '''事件轨迹子图保存''' if resultPath is None: resultPath = os.path.dirname(os.path.abspath(__file__)) subimgpath = os.path.join(resultPath, f"{event.evtname}", "subimg") imgspath = os.path.join(resultPath, f"{event.evtname}", "imgs") if not os.path.exists(subimgpath): os.makedirs(subimgpath) if not os.path.exists(imgspath): os.makedirs(imgspath) ##(2) 保存轨迹中的子图 subimgpairs = event.save_event_subimg(subimgpath) for subimgName, subimg in subimgpairs: spath = os.path.join(subimgpath, subimgName) cv2.imwrite(spath, subimg) ##(3) 保存序列图像 imgpairs = event.plot_save_image(imgspath) for imgname, img in imgpairs: spath = os.path.join(imgspath, imgname) cv2.imwrite(spath, img) ##(4) 保存轨迹散点图 img_cat = event.draw_tracks() trajpath = os.path.join(resultPath, "trajectory") if not os.path.exists(trajpath): os.makedirs(trajpath) traj_imgpath = os.path.join(trajpath, event.evtname+".png") cv2.imwrite(traj_imgpath, img_cat) def array2list(bboxes): ''' 将 bboxes 变换为 track 列表 bboxes: [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index] Return： lboxes：列表，列表中元素具有同一 track_id，x1y1x2y2 格式 [x1, y1, x2, y2, track_id, score, cls, frame_index, box_index] ''' lboxes = [] if len(bboxes)==0: return [] trackID = np.unique(bboxes[:, 4].astype(int)) track_ids = bboxes[:, 4].astype(int) for t_id in trackID: idx = np.where(track_ids == t_id)[0] box = bboxes[idx, :] lboxes.append(box) return lboxes class ShoppingEvent: def __init__(self, eventpath, stype="data"): '''stype: str, 'source', 'data', 'realtime', 共三种 ''' self.eventpath = eventpath self.evtname = str(Path(eventpath).stem) self.barcode = '' self.evtType = '' '''=========== path of image and video =========== ''' self.back_videopath = '' self.front_videopath = '' self.back_imgpaths = [] self.front_imgpaths = [] '''=========== process.data ===============================''' self.one2one = None self.one2n = None self.one2SN = None '''=========== 0/1_track.data =============================''' self.back_yolobboxes = [] self.back_yolofeats = [] self.back_trackerboxes = np.empty((0, 9), dtype=np.float64) ##和类doTracks兼容 self.back_trackerfeats = {} self.back_trackingboxes = [] self.back_trackingfeats = [] self.front_yolobboxes = [] self.front_yolofeats = [] self.front_trackerboxes = np.empty((0, 9), dtype=np.float64) ##和类doTracks兼容 self.front_trackerfeats = {} self.front_trackingboxes = [] self.front_trackingfeats = [] '''=========== 0/1_tracking_output.data ===================''' self.back_boxes = [] self.back_feats = [] self.front_boxes = [] self.front_feats = [] if stype=="data": self.from_datafile(eventpath) if stype=="realtime": self.from_realtime_datafile(eventpath) if stype=="source": self.from_source_pkl(eventpath) self.feats_select = np.empty((0, 256), dtype=np.float64) self.feats_compose = np.empty((0, 256), dtype=np.float64) self.select_feats() self.compose_feats() # if stype=="image": # self.from_image(eventpath) def kerndata(self, ShoppingDict, camtype="backCamera"): ''' camtype: str, "backCamera" or "frontCamera" ''' yoloboxes, resfeats = [], [] trackerboxes = np.empty((0, 9), dtype=np.float64) trackefeats = {} trackingboxes, trackingfeats = [], [] frameDictList = ShoppingDict[camtype]["yoloResnetTracker"] for frameDict in frameDictList: yoloboxes.append(frameDict["bboxes"]) tboxes = frameDict["tboxes"] trackefeats.update(frameDict["feats"]) trackerboxes = np.concatenate((trackerboxes, np.array(tboxes)), axis=0) Residual = ShoppingDict[camtype]["tracking"].Residual for track in Residual: trackingboxes.append(track.boxes) trackingfeats.append(track.features) kdata = (yoloboxes, resfeats, trackerboxes, trackefeats, trackingboxes, trackingfeats) tracking_out_boxes, tracking_out_feats = [], [] Confirmed = ShoppingDict[camtype]["tracking"].Confirmed for track in Confirmed: tracking_out_boxes.append(track.boxes) tracking_out_feats.append(track.features) outdata = (tracking_out_boxes, tracking_out_feats) return kdata, outdata def from_source_pkl(self, eventpath): with open(eventpath, 'rb') as f: ShoppingDict = pickle.load(f) self.eventpath = ShoppingDict["eventPath"] self.evtname = ShoppingDict["eventName"] self.barcode = ShoppingDict["barcode"] if len(ShoppingDict["one2n"]): self.one2n = ShoppingDict["one2n"] '''=========== path of image and video =========== ''' self.back_videopath = ShoppingDict["backCamera"]["videoPath"] self.front_videopath = ShoppingDict["frontCamera"]["videoPath"] self.back_imgpaths = ShoppingDict["backCamera"]["imagePaths"] self.front_imgpaths = ShoppingDict["frontCamera"]["imagePaths"] '''===========对应于 0/1_track.data =============================''' backdata, back_outdata = self.kerndata(ShoppingDict, "backCamera") frontdata, front_outdata = self.kerndata(ShoppingDict, "frontCamera") self.back_yolobboxes = backdata[0] self.back_yolofeats = backdata[1] self.back_trackerboxes = backdata[2] self.back_trackerfeats = [3] self.back_trackingboxes = [4] self.back_trackingfeats = [5] self.front_yolobboxes = frontdata[0] self.front_yolofeats = frontdata[1] self.front_trackerboxes = frontdata[2] self.front_trackerfeats = frontdata[3] self.front_trackingboxes = frontdata[4] self.front_trackingfeats = frontdata[5] '''===========对应于 0/1_tracking_output.data =============================''' self.back_boxes = back_outdata[0] self.back_feats = back_outdata[1] self.front_boxes = front_outdata[0] self.front_feats = front_outdata[1] def from_datafile(self, eventpath): evtList = self.evtname.split('_') if len(evtList)>=2 and len(evtList[-1])>=10 and evtList[-1].isdigit(): self.barcode = evtList[-1] if len(evtList)==3 and evtList[-1]== evtList[-2]: self.evtType = 'input' else: self.evtType = 'other' '''================ path of image =============''' frontImgs, frontFid = [], [] backImgs, backFid = [], [] for imgname in os.listdir(eventpath): name, ext = os.path.splitext(imgname) if ext not in IMG_FORMAT or name.find('frameId') < 0: continue if len(name.split('_')) != 3 and not name.split('_')[3].isdigit(): continue CamerType = name.split('_')[0] frameId = int(name.split('_')[3]) imgpath = os.path.join(eventpath, imgname) if CamerType == '0': backImgs.append(imgpath) backFid.append(frameId) if CamerType == '1': frontImgs.append(imgpath) frontFid.append(frameId) ## 生成依据帧 ID 排序的前后摄图像地址列表 frontIdx = np.argsort(np.array(frontFid)) backIdx = np.argsort(np.array(backFid)) self.front_imgpaths = [frontImgs[i] for i in frontIdx] self.back_imgpaths = [backImgs[i] for i in backIdx] '''================ path of video =============''' for vidname in os.listdir(eventpath): name, ext = os.path.splitext(vidname) if ext not in VID_FORMAT: continue vidpath = os.path.join(eventpath, vidname) CamerType = name.split('_')[0] if CamerType == '0': self.back_videopath = vidpath if CamerType == '1': self.front_videopath = vidpath '''================ process.data =============''' procpath = Path(eventpath).joinpath('process.data') if procpath.is_file(): SimiDict = read_similar(procpath) self.one2one = SimiDict['one2one'] self.one2n = SimiDict['one2n'] self.one2SN = SimiDict['one2SN'] '''=========== 0/1_track.data & 0/1_tracking_output.data =======''' for dataname in os.listdir(eventpath): datapath = os.path.join(eventpath, dataname) if not os.path.isfile(datapath): continue CamerType = dataname.split('_')[0] '''========== 0/1_track.data ==========''' if dataname.find("_track.data")>0: bboxes, ffeats, trackerboxes, trackerfeats, trackingboxes, trackingfeats = extract_data(datapath) if CamerType == '0': self.back_yolobboxes = bboxes self.back_yolofeats = ffeats self.back_trackerboxes = trackerboxes self.back_trackerfeats = trackerfeats self.back_trackingboxes = trackingboxes self.back_trackingfeats = trackingfeats if CamerType == '1': self.front_yolobboxes = bboxes self.front_yolofeats = ffeats self.front_trackerboxes = trackerboxes self.front_trackerfeats = trackerfeats self.front_trackingboxes = trackingboxes self.front_trackingfeats = trackingfeats '''========== 0/1_tracking_output.data ==========''' if dataname.find("_tracking_output.data")>0: tracking_output_boxes, tracking_output_feats = read_tracking_output(datapath) if CamerType == '0': self.back_boxes = tracking_output_boxes self.back_feats = tracking_output_feats elif CamerType == '1': self.front_boxes = tracking_output_boxes self.front_feats = tracking_output_feats def from_realtime_datafile(self, eventpath): # evtList = self.evtname.split('_') # if len(evtList)>=2 and len(evtList[-1])>=10 and evtList[-1].isdigit(): # self.barcode = evtList[-1] # if len(evtList)==3 and evtList[-1]== evtList[-2]: # self.evtType = 'input' # else: # self.evtType = 'other' '''================ path of video =============''' for vidname in os.listdir(eventpath): name, ext = os.path.splitext(vidname) if ext not in VID_FORMAT: continue vidpath = os.path.join(eventpath, vidname) CamerType = name.split('_')[0] if CamerType == '0': self.back_videopath = vidpath if CamerType == '1': self.front_videopath = vidpath '''================ process.data =============''' procpath = Path(eventpath).joinpath('process.data') if procpath.is_file(): SimiDict = read_similar(procpath) self.one2one = SimiDict['one2one'] self.one2n = SimiDict['one2n'] self.one2SN = SimiDict['one2SN'] '''=========== 0/1_track.data & 0/1_tracking_output.data =======''' for dataname in os.listdir(eventpath): datapath = os.path.join(eventpath, dataname) if not os.path.isfile(datapath): continue CamerType = dataname.split('_')[0] '''========== 0/1_track.data ==========''' if dataname.find("_track.data")>0: trackerboxes, trackerfeats = extract_data_realtime(datapath) if CamerType == '0': self.back_trackerboxes = trackerboxes self.back_trackerfeats = trackerfeats if CamerType == '1': self.front_trackerboxes = trackerboxes self.front_trackerfeats = trackerfeats '''========== 0/1_tracking_output.data ==========''' if dataname.find("_tracking_output.data")>0: trackingboxes, trackingfeats, tracking_outboxes, tracking_outfeats = read_tracking_output_realtime(datapath) if CamerType == '0': self.back_trackingboxes = trackingboxes self.back_trackingfeats = trackingfeats self.back_boxes = tracking_outboxes self.back_feats = tracking_outfeats elif CamerType == '1': self.front_trackingboxes = trackingboxes self.front_trackingfeats = trackingfeats self.front_boxes = tracking_outboxes self.front_feats = tracking_outfeats def compose_feats(self): '''事件的特征集成''' feats_compose = np.empty((0, 256), dtype=np.float64) if len(self.front_feats): for feat in self.front_feats: feats_compose = np.concatenate((feats_compose, feat), axis=0) if len(self.back_feats): for feat in self.back_feats: feats_compose = np.concatenate((feats_compose, feat), axis=0) self.feats_compose = feats_compose def select_feats(self): '''事件的特征选择''' if len(self.front_feats): self.feats_select = self.front_feats[0] elif len(self.back_feats): self.feats_select = self.back_feats[0] def plot_save_image(self, savepath): def array2list(bboxes): '''[x1, y1, x2, y2, track_id, score, cls, frame_index, box_index]''' frame_ids = bboxes[:, 7].astype(int) fID = np.unique(bboxes[:, 7].astype(int)) fboxes = [] for f_id in fID: idx = np.where(frame_ids==f_id)[0] box = bboxes[idx, :] fboxes.append((f_id, box)) return fboxes imgpairs = [] cameras = ('front', 'back') for camera in cameras: if camera == 'front': boxes = self.front_trackerboxes imgpaths = self.front_imgpaths else: boxes = self.back_trackerboxes imgpaths = self.back_imgpaths fboxes = array2list(boxes) for fid, fbox in fboxes: imgpath = imgpaths[int(fid-1)] image = cv2.imread(imgpath) annotator = Annotator(image.copy(), line_width=2) for i, box in enumerate(fbox): x1, y1, x2, y2, tid, score, cls, fid, bid = box label = f'{int(tid), int(cls)}' if tid >=0 and cls==0: color = colors(int(cls), True) elif tid >=0 and cls!=0: color = colors(int(tid), True) else: color = colors(19, True) # 19为调色板的最后一个元素 xyxy = (x1/2, y1/2, x2/2, y2/2) annotator.box_label(xyxy, label, color=color) im0 = annotator.result() imgpairs.append((Path(imgpath).name, im0)) # spath = os.path.join(savepath, Path(imgpath).name) # cv2.imwrite(spath, im0) return imgpairs def save_event_subimg(self, savepath): ''' 功能: 保存一次购物事件的轨迹子图 9 items: barcode, type, filepath, back_imgpaths, front_imgpaths, back_boxes, front_boxes, back_feats, front_feats, feats_compose, feats_select 子图保存次序：先前摄、后后摄，以 k 为编号，和 "feats_compose" 中次序相同 ''' imgpairs = [] cameras = ('front', 'back') for camera in cameras: boxes = np.empty((0, 9), dtype=np.float64) ##和类doTracks兼容 if camera == 'front': for b in self.front_boxes: boxes = np.concatenate((boxes, b), axis=0) imgpaths = self.front_imgpaths else: for b in self.back_boxes: boxes = np.concatenate((boxes, b), axis=0) imgpaths = self.back_imgpaths for i, box in enumerate(boxes): x1, y1, x2, y2, tid, score, cls, fid, bid = box imgpath = imgpaths[int(fid-1)] image = cv2.imread(imgpath) subimg = image[int(y1/2):int(y2/2), int(x1/2):int(x2/2), :] camerType, timeTamp, _, frameID = os.path.basename(imgpath).split('.')[0].split('_') subimgName = f"cam{camerType}_{i}_tid{int(tid)}_fid({int(fid)}, {frameID}).png" imgpairs.append((subimgName, subimg)) # spath = os.path.join(savepath, subimgName) # cv2.imwrite(spath, subimg) return imgpairs # basename = os.path.basename(event['filepath']) print(f"Image saved: {os.path.basename(self.eventpath)}") def draw_tracks(self): front_edge = cv2.imread(r"D:\DetectTracking\tracking\shopcart\cart_tempt\board_ftmp_line.png") back_edge = cv2.imread(r"D:\DetectTracking\tracking\shopcart\cart_tempt\edgeline.png") front_trackerboxes = array2list(self.front_trackerboxes) back_trackerboxes = array2list(self.back_trackerboxes) # img1, img2 = edgeline.copy(), edgeline.copy() img1 = drawTrack(front_trackerboxes, front_edge.copy()) img2 = drawTrack(self.front_trackingboxes, front_edge.copy()) img3 = drawTrack(back_trackerboxes, back_edge.copy()) img4 = drawTrack(self.back_trackingboxes, back_edge.copy()) imgcat1 = np.concatenate((img1, img2), axis = 1) H, W = imgcat1.shape[:2] cv2.line(imgcat1, (int(W/2), 0), (int(W/2), H), (128, 255, 128), 2) imgcat2 = np.concatenate((img3, img4), axis = 1) H, W = imgcat2.shape[:2] cv2.line(imgcat2, (int(W/2), 0), (int(W/2), H), (128, 255, 128), 2) illus = [imgcat1, imgcat2] 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) return img_cat def main(): # pklpath = r"D:\DetectTracking\evtresult\images2\ShoppingDict.pkl" # evt = ShoppingEvent(pklpath, stype='source') evtpath = r"\\192.168.1.28\share\测试视频数据以及日志\算法全流程测试\202412\images\20241209-160248-08edd5f6-1806-45ad-babf-7a4dd11cea60_6973226721445" evt = ShoppingEvent(evtpath, stype='data') img_cat = evt.draw_tracks() cv2.imwrite("a.png", img_cat) if __name__ == "__main__": main() # main1()