回传数据解析，兼容v5和v10

tracking/tracking_test.py

@@ -0,0 +1,176 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Sep 20 17:33:00 2023
+
+@author: ym
+"""
+
+import cv2
+import os
+import numpy as np
+import time
+import pickle
+import matplotlib.pyplot as plt
+import pandas as pd 
+from scipy.spatial.distance import cdist
+
+from pathlib import Path
+# ================= using for import ultralytics
+import sys
+sys.path.append(r"D:\DetectTracking")
+from utils.gen import Profile
+
+from dotrack.dotracks_back import doBackTracks
+from dotrack.dotracks_front import doFrontTracks 
+from utils.drawtracks import draw5points, drawTrack, drawtracefeat, plot_frameID_y2, drawFeatures, draw_all_trajectories
+
+from utils.read_data import extract_data_realtime, read_tracking_output_realtime
+
+# from datetime import datetime
+# from utils.proBoxes import boxes_add_fid
+# from utils.plotting import boxing_img  #, Annotator, colors,
+# from utils import Boxes, IterableSimpleNamespace, yaml_load
+# from trackers import BOTSORT, BYTETracker
+# from utils.mergetrack import track_equal_track
+# from utils.basetrack import MoveState, ShoppingCart, doTracks
+
+
+def detect_start_end(bboxes, features_dict, filename):
+    
+    boxes = np.empty(shape=(0, 9), dtype = np.float)
+    if filename.find("back") >= 0:
+        vts = doBackTracks(bboxes, features_dict)
+        vtx = [t for t in vts if t.cls != 0]     
+        for track in vtx:
+            if track.moving_index.size:
+                boxes = np.concatenate((boxes, track.moving_index), axis=0)
+                
+    elif filename.find("front") >= 0:
+        vts = doFrontTracks(bboxes, features_dict)
+        vtx = [t for t in vts if t.cls != 0]
+        
+        for track in vtx:
+            for start, end in track.dynamic_y2:
+                boxes = np.concatenate((boxes, track.boxes[start:end+1, :]), axis=0)    
+            for start, end in track.dynamic_y1:
+                boxes = np.concatenate((boxes, track.boxes[start:end+1, :]), axis=0)
+    
+    start = np.min(boxes[:, 0])
+    end = np.max(boxes[:, 1])
+    
+    if start > 5:
+        start = start - 5
+    else:
+        start = 0
+    
+    return start, end
+    
+def save_subimgs(vts, file, TracksDict):
+    imgdir = Path(f'./result/imgs/{file}')
+    if not imgdir.exists():
+        imgdir.mkdir(parents=True, exist_ok=True)
+    
+    for i, track in enumerate(vts.Residual):  
+        boxes = track.boxes
+        for ii in range(len(boxes)):
+            tid, fid, bid = int(boxes[ii, 4]), int(boxes[ii, 7]), int(boxes[ii, 8])
+            
+            img = TracksDict[f"frame_{fid}"]["imgs"][bid]
+            # feat = TracksDict[f"frame_{fid}"]["feats"][bid]
+            # box = TracksDict[f"frame_{fid}"]["boxes"][bid]
+            
+            cv2.imwrite(str(imgdir) + f"/{tid}_{fid}_{bid}.png", img)
+
+def have_tracked():
+    # trackdict = r'./data/trackdicts'
+    trackdict = r'D:\全实时\202502\20250228-152846-438_6901668934727_6901668934727'
+    
+    bboxes, alltracks = [], []
+    k = 0
+    gt = Profile()
+    for filename in os.listdir(trackdict):
+        # filename = '153112511_0_seek_105.pkl'
+        
+        file, ext = os.path.splitext(filename)
+        filepath = os.path.join(trackdict, filename)
+        
+        if file.split('_')[0]=='0' or file.find("back") >= 0:
+            CamerType = "back"
+        if file.split('_')[0]=='1' or file.find("front") >= 0:
+            CamerType = "front"
+        
+        ## 1. 加载 tracker 输出的 pickle 文件
+        if ext in ['.pkl', '.pickle']:
+            filepath = os.path.join(trackdict, filename)
+            TracksDict = np.load(filepath, allow_pickle=True)
+            bboxes = TracksDict['TrackBoxes']
+        
+        ## 2. 加载 data 文件    
+        if filename.find('tracker.data')>0:
+            bboxes, TracksDict = extract_data_realtime(filepath)
+            
+        if len(bboxes)==0:
+            continue
+
+        with gt:
+            if CamerType == "front":
+                vts = doFrontTracks(bboxes, TracksDict)
+                
+                Intrude = vts.isintrude()
+                
+                vts.classify()
+                
+                save_subimgs(vts, file, TracksDict)
+
+                plt = plot_frameID_y2(vts)
+                
+                savedir = save_dir.joinpath(f'{file}_y2.png') 
+                
+                plt.savefig(savedir)  
+                plt.close()
+                
+                edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
+                img_tracking = draw_all_trajectories(vts, edgeline, save_dir, file, draw5p=True)
+                
+            if CamerType == "back":
+                vts = doBackTracks(bboxes, TracksDict)
+                
+                Intrude = vts.isintrude()
+                
+                vts.classify()
+                alltracks.append(vts)
+                
+                save_subimgs(vts, file, TracksDict)
+                edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
+                img_tracking = draw_all_trajectories(vts, edgeline, save_dir, file)
+                
+            trackpath = save_dir.joinpath(f'{file}.png')
+            cv2.imwrite(str(trackpath), img_tracking)
+        print(file+f" need time: {gt.dt:.2f}s")
+
+        k += 1
+        if k==1:
+            break
+    
+    if len(alltracks):
+        drawFeatures(alltracks, save_dir)
+         
+    
+    
+if __name__ == "__main__": 
+    # now = datetime.now()
+    # time_string = now.strftime("%Y%m%d%H%M%S")[:8]
+    
+    save_dir = Path('./result/tracks')
+    if not save_dir.exists():
+        save_dir.mkdir(parents=True, exist_ok=True)
+
+    have_tracked()    
+   
+
+    
+    
+    
+    
+    
+