update

ytracking/tracking/have_tracking.py

@@ -0,0 +1,254 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Sep 20 17:33:00 2023
+
+@author: ym
+"""
+import sys
+
+import cv2
+import os
+import numpy as np
+import time
+import pickle
+import matplotlib.pyplot as plt
+from scipy.spatial.distance import cdist
+
+from pathlib import Path
+from datetime import datetime
+
+# ================= using for import ultralytics
+
+# sys.path.append(r"D:\DeepLearning\yolov5_track")
+
+# from utils.proBoxes import boxes_add_fid
+from ytracking.tracking.utils.plotting import boxing_img  # , Annotator, colors,
+from ytracking.tracking.utils.gen import Profile
+from ytracking.tracking.utils.drawtracks import draw5points, drawTrack, drawtracefeat, plot_frameID_y2, drawFeatures, \
+    draw_all_trajectories
+from ytracking.tracking.utils import Boxes, IterableSimpleNamespace, yaml_load
+from ytracking.tracking.trackers import BOTSORT, BYTETracker
+
+sys.path.append("ytracking/tracking/")
+from dotrack.dotracks_back import doBackTracks
+from dotrack.dotracks_front import doFrontTracks
+
+
+# from utils.mergetrack import track_equal_track
+# from utils.basetrack import MoveState, ShoppingCart, doTracks
+
+def init_tracker(tracker_yaml=None, bs=1):
+    """
+    Initialize tracker for object tracking during prediction.
+    """
+    TRACKER_MAP = {'bytetrack': BYTETracker, 'botsort': BOTSORT}
+    cfg = IterableSimpleNamespace(**yaml_load(tracker_yaml))
+
+    tracker = TRACKER_MAP[cfg.tracker_type](args=cfg, frame_rate=30)
+
+    return tracker
+
+
+def have_tracked_front():
+    '''前摄轨迹处理。已执行跟踪处理，只对跟踪结果进行分析'''
+
+    featdir = r"./data/trackfeats"
+    npydir = r"./data/tracks"
+    k = 0
+    gt = Profile()
+    for filename in os.listdir(npydir):
+        # filename = "084501222314_20240108-143651_front.npy"
+        if not filename.find("front") >= 0: continue
+
+        file, ext = os.path.splitext(filename)
+        fpath = os.path.join(npydir, filename)
+        featpath = os.path.join(featdir, file + '.pkl')
+
+        bboxes = np.load(fpath)
+        features_dict = np.load(featpath, allow_pickle=True)
+
+        with gt:
+            vts = doFrontTracks(bboxes, features_dict)
+            vts.classify()
+
+            plt = plot_frameID_y2(vts)
+            plt.savefig(f'./result/{file}_y2.png')
+            plt.close()
+
+        print(file + f" need time: {gt.dt:.2f}s")
+
+        # edgeline = cv2.imread("./shopcart/cart_tempt/board_ftmp_line.png")
+        # draw_all_trajectories(vts, edgeline, save_dir, filename)
+
+        # k += 1
+        # if k == 1:
+        #     break
+
+
+def have_tracked_back():
+    '''后摄轨迹处理。已执行跟踪处理，只对跟踪结果进行分析'''
+    featdir = r"./data/trackfeats"
+    npydir = r"./data/tracks"
+    k = 0
+    alltracks = []
+    gt = Profile()
+    for filename in os.listdir(npydir):
+        # filename = "084501222314_20240108-143656_back.npy"     # "加购_55.npy"
+        if not filename.find("back") >= 0: continue
+
+        t1 = time.time()
+
+        file, ext = os.path.splitext(filename)
+        fpath = os.path.join(npydir, filename)
+        featpath = os.path.join(featdir, file + '.pkl')
+
+        # try:   
+        bboxes = np.load(fpath)
+        features_dict = np.load(featpath, allow_pickle=True)
+
+        with gt:
+            vts = doBackTracks(bboxes, features_dict)
+            vts.classify()
+            # vts.merge_tracks()
+
+        print(file + f" need time: {gt.dt:.2f}s")
+        edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
+        draw_all_trajectories(vts, edgeline, save_dir, filename)
+        alltracks.append(vts)
+
+        # except Exception as e:
+        #     # print(str(e))
+        #     pass
+        # print(file+" need time: {:.2f}s".format(time.time()-t1))
+        k += 1
+        if k == 1:
+            break
+
+    if len(alltracks):
+        drawFeatures(alltracks, save_dir)
+
+
+def tracking(vboxes):
+    tracker_yaml = r"./trackers/cfg/botsort.yaml"
+    tracker = init_tracker(tracker_yaml)
+    tboxes = []
+    images = []
+
+    track_boxes = np.empty((0, 9), dtype=np.float32)
+    features_dict = {}
+
+    '''==================== 执行跟踪处理 ======================='''
+    for det, img, frame in vboxes:
+        # 需要根据frame_id重排序
+        det_tracking = Boxes(det).cpu().numpy()
+        H, W = img.shape[:2]
+        imgs = []
+        for d in range(np.size(det, 0)):
+            tlbr = det[d, :4].astype(np.int_)
+            tlbr[0] = max(0, tlbr[0])
+            tlbr[1] = max(0, tlbr[1])
+            tlbr[2] = min(W - 1, tlbr[2])
+            tlbr[3] = min(H - 1, tlbr[3])
+            patch = img[tlbr[1]:tlbr[3], tlbr[0]:tlbr[2], :]
+            patch = patch[:, :, ::-1]  # 原程序用PIL.Image读取，是RGB，Opencv是BGR
+            imgs.append(patch)
+
+        tracks = tracker.update(det_tracking, imgs)
+
+        if len(tracks):
+            track_boxes = np.concatenate([track_boxes, tracks], axis=0)
+            feat_dict = {int(x.idx): x.curr_feat for x in tracker.tracked_stracks if x.is_activated}
+            frame_id = tracks[0, 7]
+            features_dict.update({int(frame_id): feat_dict})
+
+            # det = tracks[:, :-1]
+            # tboxes.append((det, frame))
+
+            imgx = boxing_img(tracks, img)
+            images.append((imgx, frame))
+
+    # bboxes = boxes_add_fid(tboxes)
+
+    vts = doBackTracks(track_boxes, features_dict)
+    vts.classify()
+
+    return vts, images
+
+
+def do_tracking():
+    pkldir = r"./data/boxes_imgs"
+    k = 0
+    save_result = True
+    alltracks = []
+    gt = Profile()
+    for filename in os.listdir(pkldir):
+        filename = "加购_18.pkl"
+        file, _ = os.path.splitext(filename)
+        vboxes = []
+        ##================================ load the detection data
+        with open(pkldir + f'/{filename}', 'rb') as f:
+            vboxes = pickle.load(f)
+        assert len(vboxes) > 0
+
+        with gt:
+            vts, images = tracking(vboxes)
+            alltracks.append(vts)
+
+        print(file + f" need time: {gt.dt * 1E3:.1f}ms")
+        ##================================ save images, video, track-trajectory
+        if save_result == True:
+            curdir = imgdir.joinpath(file)
+            if not curdir.exists():
+                curdir.mkdir(parents=True, exist_ok=True)
+            vidpath = str(curdir.joinpath(file).with_suffix('.mp4'))
+
+            fps, w, h = 30, images[0][0].shape[1], images[0][0].shape[0]
+            vidwriter = cv2.VideoWriter(vidpath, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
+            for img, frame in images:
+                imgpath = curdir.joinpath(file + f"_{frame}.png")
+                cv2.imwrite(imgpath, img)
+                vidwriter.write(img)
+            vidwriter.release()
+            edgeline = cv2.imread("./shopcart/cart_tempt/edgeline.png")
+            draw_all_trajectories(vts, edgeline, save_dir, filename)
+
+        k += 1
+        if k == 1:
+            break
+
+    drawFeatures(alltracks, save_dir)
+
+
+def have_tracked(bboxes, features_dict, camera_id):
+    if camera_id == '0':
+        vts = doBackTracks(bboxes, features_dict)
+        vts.classify()
+    elif camera_id == '1':
+        vts = doFrontTracks(bboxes, features_dict)
+        vts.classify()
+    else:
+        raise ValueError("have no camera_id")
+    return vts
+
+
+if __name__ == "__main__":
+    now = datetime.now()
+    time_string = now.strftime("%Y%m%d%H%M%S")[:8]
+
+    # save_dir = Path(f'./result/{time_string}_traj/')
+    # if not save_dir.exists():
+    #     save_dir.mkdir(parents=True, exist_ok=True)
+
+    save_dir = Path(f'./result/')
+
+    mode = "merge"  ## "merge": 已完成跟踪处理, "other": 未执行跟踪处理
+    if mode == "merge":
+        # have_tracked_back()
+        have_tracked_front()
+    else:
+        '''执行do_tracking()函数时视频和图像存储位置'''
+        imgdir = Path(f'./result/{time_string}_imgs/')
+        if not imgdir.exists():
+            imgdir.mkdir(parents=True, exist_ok=True)
+
+        do_tracking()