This commit is contained in:
王庆刚
2024-11-25 18:05:08 +08:00
parent c47894ddc0
commit 8bbee310ba
109 changed files with 1003 additions and 305 deletions

View File

@ -10,8 +10,15 @@ from .utils import matching
# from .utils.gmc import GMC
from .utils.kalman_filter import KalmanFilterXYWH
from .reid.reid_interface import ReIDInterface
from .reid.config import config
# from .reid.reid_interface import ReIDInterface
# from .reid.config import config
from contrast.feat_extract.inference import FeatsInterface
from contrast.feat_extract.config import config as conf
class BOTrack(STrack):
shared_kalman = KalmanFilterXYWH()
@ -111,7 +118,9 @@ class BOTSORT(BYTETracker):
if args.with_reid:
# Haven't supported BoT-SORT(reid) yet
self.encoder = ReIDInterface(config)
# self.encoder = ReIDInterface(config)
self.encoder = FeatsInterface(conf)
# self.gmc = GMC(method=args.gmc_method) # commented by WQG
@ -119,13 +128,13 @@ class BOTSORT(BYTETracker):
"""Returns an instance of KalmanFilterXYWH for object tracking."""
return KalmanFilterXYWH()
def init_track(self, dets, scores, cls, imgs, features_keep):
def init_track(self, dets, scores, cls, image, features_keep):
"""Initialize track with detections, scores, and classes."""
if len(dets) == 0:
return []
if self.args.with_reid and self.encoder is not None:
if features_keep is None:
features_keep = self.encoder.inference(imgs, dets)
imgs, features_keep = self.encoder.inference(image, dets)
return [BOTrack(xyxy, s, c, f) for (xyxy, s, c, f) in zip(dets, scores, cls, features_keep)] # detections
else: