# -*- coding: utf-8 -*- """ Created on Mon Oct 14 10:01:24 2024 @author: ym """ import numpy as np import cv2 from scipy.spatial.distance import cdist class TrackFrag: def __init__(self, boxes, imgshape=(1280, 1024)): self.boxes = boxes self.cls = int(boxes[0, 6]) self.tid = int(boxes[0, 4]) self.imgshape = imgshape '''轨迹的持续时间,以帧ID表征, (或考虑用时间戳)''' self.during = (np.min(boxes[:, 7]), np.max(boxes[:, 7])) self.groups = [set(np.unique(boxes[:, 7].astype(int)))] # '''5个关键点(中心点、左上点、右上点、左下点、右下点 )坐标''' self.isCornpoint = self.is_cornpoint(10) self.compute_cornpoints() def is_cornpoint(self, edge=10): isleft = min(self.boxes[:, 0]) < edge istop = min(self.boxes[:, 1]) < edge isright = max(self.boxes[:, 2]) > self.imgshape[0] - edge isbottom = max(self.boxes[:, 3]) > self.imgshape[0] - edge isCornpoint = isbottom or istop or isleft or isright return isCornpoint def compute_cornpoints(self): ''' cornpoints 共10项,分别是个点的坐标值(x, y) (center, top_left, top_right, bottom_left, bottom_right) ''' boxes = self.boxes cornpoints = np.zeros((len(boxes), 10)) cornpoints[:,0] = (boxes[:, 0] + boxes[:, 2]) / 2 cornpoints[:,1] = (boxes[:, 1] + boxes[:, 3]) / 2 cornpoints[:,2], cornpoints[:,3] = boxes[:, 0], boxes[:, 1] cornpoints[:,4], cornpoints[:,5] = boxes[:, 2], boxes[:, 1] cornpoints[:,6], cornpoints[:,7] = boxes[:, 0], boxes[:, 3] cornpoints[:,8], cornpoints[:,9] = boxes[:, 2], boxes[:, 3] trajdist = [] for k in range(5): X = cornpoints[:, 2*k:2*(k+1)] trajdist.append(np.max(cdist(X, X))) idx = trajdist.index(min(trajdist)) self.trajdist_min = trajdist[idx] self.cornpoints = cornpoints def update_groups(self, THRESH=18): ''' 对 self.groups 重新赋值 ''' boxes = self.boxes nbox = len(boxes) X = np.zeros((len(boxes), 2)) X[:,0] = (boxes[:, 0] + boxes[:, 2]) / 2 X[:,1] = (boxes[:, 1] + boxes[:, 3]) / 2 dist2 = cdist(X, X) # label = np.zeros(nbox, dtype=np.int) marked, groups = set(), [] for k in range(nbox): if k in marked: continue group = set() dt = dist2[k, :] idx = np.where(dt < THRESH)[0] if len(idx) == 1: groups.append({k}) marked.add(k) continue '''初始近邻样本点集合, 并移除当前点''' seeds = set(idx) seeds.remove(k) group.add(k) marked.add(k) while len(seeds) !=0: pt = seeds.pop() dt = dist2[pt, :] seed = set(np.where(dt < THRESH)[0]) seed.remove(pt) seed.difference_update(marked) seeds.update(seed) group.add(pt) marked.add(pt) groups.append(group) self.groups = groups def jump_boxes(self): gpboxes = [] for group in self.groups: box = self.boxes[list(group), :] gpboxes.append(box) return gpboxes def is_moving(self): if len(self.groups)>=3: return True return False def is_static(self, THRESH=50): box1 = self.boxes[0, :4] box2 = self.boxes[-1, :4] ''' 第1帧、最后一帧 boxes 四个角点间的距离 ''' ptd = box2 - box1 ptd1 = np.linalg.norm((ptd[0], ptd[1])) ptd2 = np.linalg.norm((ptd[2], ptd[1])) ptd3 = np.linalg.norm((ptd[0], ptd[3])) ptd4 = np.linalg.norm((ptd[2], ptd[3])) condt1 = ptd1=3] def draw(self): pass def array2list(self): ''' 将 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] ''' track_ids = self.bboxes[:, 4].astype(int) lboxes = [] for t_id in self.trackIDs: # print(f"The ID is: {t_id}") idx = np.where(track_ids == t_id)[0] box = self.bboxes[idx, :] lboxes.append(box) return lboxes @staticmethod def sub_tracks(tlista, tlistb): track_ids_b = {t.tid for t in tlistb} return [t for t in tlista if t.tid not in track_ids_b]