148 lines
4.3 KiB
Python
148 lines
4.3 KiB
Python
# -*- 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 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'
|
|
alltracks = []
|
|
k = 0
|
|
gt = Profile()
|
|
for filename in os.listdir(trackdict):
|
|
# filename = 'test_20240402-173935_6920152400975_back_174037372.pkl'
|
|
# filename = '加购_91.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, TracksDict)
|
|
|
|
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, TracksDict)
|
|
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()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|