|
|
|
|
|
|
|
|
|
|
|
import numpy as np |
|
|
from .visualization import find_peak_box_improved, convert_grid_to_xy |
|
|
|
|
|
class TrackedObject: |
|
|
def __init__(self): |
|
|
self.last_step=0;self.last_pos=[0,0];self.historical_pos=[];self.historical_steps=[];self.historical_features=[] |
|
|
def update(self,step,object_info): |
|
|
self.last_step=step;self.last_pos=object_info[:2];self.feature=object_info[2:];self.historical_pos.append(self.last_pos);self.historical_steps.append(step);self.historical_features.append(self.feature) |
|
|
|
|
|
class Tracker: |
|
|
def __init__(self,frequency=10):self.tracks=[];self.alive_ids=[];self.frequency=frequency |
|
|
def update_and_predict(self,det_data,pos,theta,step,merge_percent=0.4): |
|
|
detected_objects=find_peak_box_improved(det_data);objects_info=[];R=np.array([[np.cos(-theta),-np.sin(-theta)],[np.sin(-theta),np.cos(-theta)]]);box_ids_for_update=[] |
|
|
for obj in detected_objects: |
|
|
i,j=obj['coords'];obj_data=obj['raw_data'];center_y,center_x=convert_grid_to_xy(i,j);center_x+=obj_data[1];center_y+=obj_data[2];loc=R.T.dot(np.array([center_x,center_y])) |
|
|
objects_info.append([loc[0]+pos[0],loc[1]+pos[1],obj_data[1:]]);box_ids_for_update.append((i,j)) |
|
|
self._update(objects_info,step);return det_data |
|
|
def _update(self,objects_info,step): |
|
|
if not self.tracks: |
|
|
for info in objects_info:self.tracks.append(TrackedObject());self.tracks[-1].update(step,info) |
|
|
else: |
|
|
to_ids=self._match(objects_info) |
|
|
for i,to_id in enumerate(to_ids): |
|
|
if to_id==-1:self.tracks.append(TrackedObject());to_id=len(self.tracks)-1 |
|
|
self.tracks[to_id].update(step,objects_info[i]) |
|
|
self.alive_ids=[i for i,track in enumerate(self.tracks) if track.last_step>step-6] |
|
|
def _match(self,objects_info): |
|
|
results=[];matched_ids=set() |
|
|
for info in objects_info: |
|
|
min_id,min_error=-1,float('inf');pos_x,pos_y=info[:2] |
|
|
for _id in self.alive_ids: |
|
|
if _id in matched_ids:continue |
|
|
track_pos=self.tracks[_id].last_pos;distance=np.sqrt((track_pos[0]-pos_x)**2+(track_pos[1]-pos_y)**2) |
|
|
if distance<2.0 and distance<min_error:min_error=distance;min_id=_id |
|
|
results.append(min_id) |
|
|
if min_id!=-1:matched_ids.add(min_id) |
|
|
return results |
|
|
|
