defon_predict_start(predictor,persist=False):""" Initialize trackers for object tracking during prediction. Args: predictor (object): The predictor object to initialize trackers for. persist (bool, optional): Whether to persist the trackers if they already exist. Defaults to False. Raises: AssertionError: If the tracker_type is not 'bytetrack' or 'botsort'. """ifhasattr(predictor,'trackers')andpersist:returntracker=check_yaml(predictor.args.tracker)cfg=IterableSimpleNamespace(**yaml_load(tracker))assertcfg.tracker_typein['bytetrack','botsort'], \
f"Only support 'bytetrack' and 'botsort' for now, but got '{cfg.tracker_type}'"trackers=[]for_inrange(predictor.dataset.bs):tracker=TRACKER_MAP[cfg.tracker_type](args=cfg,frame_rate=30)trackers.append(tracker)predictor.trackers=trackers
on_predict_postprocess_end
Postprocess detected boxes and update with object tracking.
defon_predict_postprocess_end(predictor):"""Postprocess detected boxes and update with object tracking."""bs=predictor.dataset.bsim0s=predictor.batch[1]foriinrange(bs):det=predictor.results[i].boxes.cpu().numpy()iflen(det)==0:continuetracks=predictor.trackers[i].update(det,im0s[i])iflen(tracks)==0:continueidx=tracks[:,-1].astype(int)predictor.results[i]=predictor.results[i][idx]predictor.results[i].update(boxes=torch.as_tensor(tracks[:,:-1]))
register_tracker
Register tracking callbacks to the model for object tracking during prediction.
Parameters:
Name
Type
Description
Default
model
object
The model object to register tracking callbacks for.
required
persist
bool
Whether to persist the trackers if they already exist.
defregister_tracker(model,persist):""" Register tracking callbacks to the model for object tracking during prediction. Args: model (object): The model object to register tracking callbacks for. persist (bool): Whether to persist the trackers if they already exist. """model.add_callback('on_predict_start',partial(on_predict_start,persist=persist))model.add_callback('on_predict_postprocess_end',on_predict_postprocess_end)
Created 2023-04-16, Updated 2023-05-17 Authors: Glenn Jocher (3)