initial project version!

tracking/test_tracking.py

@@ -0,0 +1,149 @@
+# -*- 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:\yolov5track")
+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 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):
+    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 = vts.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'
+    alltracks = []
+    k = 0
+    gt = Profile()
+    for filename in os.listdir(trackdict):
+        filename = 'test_20240402-173935_6920152400975_front_174037379.pkl'
+        file, ext = os.path.splitext(filename)
+        filepath = os.path.join(trackdict, filename)
+        TracksDict = np.load(filepath, allow_pickle=True)
+
+        bboxes = TracksDict['TrackBoxes']
+        with gt:
+            if filename.find("front") >= 0:
+                vts = doFrontTracks(bboxes, TracksDict)
+                vts.classify()
+                
+                save_subimgs(vts, file)
+
+                plt = plot_frameID_y2(vts)
+                
+                savedir = save_dir.joinpath(f'{file}_y2.png') 
+                
+                plt.savefig(savedir)  
+                plt.close()
+            else:
+                vts = doBackTracks(bboxes, TracksDict)
+                vts.classify()
+                alltracks.append(vts)
+                
+                save_subimgs(vts, file)
+                edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
+                draw_all_trajectories(vts, edgeline, save_dir, filename)
+        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()    
+   
+
+    
+    
+    
+    
+    
+