bytetrack.py 2.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879
  1. # Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. from __future__ import division
  16. from __future__ import print_function
  17. from paddlex.ppdet.core.workspace import register, create
  18. from .meta_arch import BaseArch
  19. __all__ = ['ByteTrack']
  20. @register
  21. class ByteTrack(BaseArch):
  22. """
  23. ByteTrack network, see https://arxiv.org/abs/2110.06864
  24. Args:
  25. detector (object): detector model instance
  26. reid (object): reid model instance, default None
  27. tracker (object): tracker instance
  28. """
  29. __category__ = 'architecture'
  30. def __init__(self, detector='YOLOX', reid=None, tracker='JDETracker'):
  31. super(ByteTrack, self).__init__()
  32. self.detector = detector
  33. self.reid = reid
  34. self.tracker = tracker
  35. @classmethod
  36. def from_config(cls, cfg, *args, **kwargs):
  37. detector = create(cfg['detector'])
  38. if cfg['reid'] != 'None':
  39. reid = create(cfg['reid'])
  40. else:
  41. reid = None
  42. tracker = create(cfg['tracker'])
  43. return {
  44. "detector": detector,
  45. "reid": reid,
  46. "tracker": tracker,
  47. }
  48. def _forward(self):
  49. det_outs = self.detector(self.inputs)
  50. if self.training:
  51. return det_outs
  52. else:
  53. if self.reid is not None:
  54. assert 'crops' in self.inputs
  55. crops = self.inputs['crops']
  56. pred_embs = self.reid(crops)
  57. else:
  58. pred_embs = None
  59. det_outs['embeddings'] = pred_embs
  60. return det_outs
  61. def get_loss(self):
  62. return self._forward()
  63. def get_pred(self):
  64. return self._forward()