mirror of https://github.com/JDAI-CV/fast-reid.git
25 lines
652 B
Python
25 lines
652 B
Python
|
# encoding: utf-8
|
||
|
|
"""
|
||
|
|
@author: liaoxingyu
|
||
|
|
@contact: sherlockliao01@gmail.com
|
||
|
|
"""
|
||
|
|
|
||
|
|
from ...utils.registry import Registry
|
||
|
|
|
||
|
|
REID_HEADS_REGISTRY = Registry("REID_HEADS")
|
||
|
|
REID_HEADS_REGISTRY.__doc__ = """
|
||
|
|
Registry for ROI heads in a generalized R-CNN model.
|
||
|
|
ROIHeads take feature maps and region proposals, and
|
||
|
|
perform per-region computation.
|
||
|
|
The registered object will be called with `obj(cfg, input_shape)`.
|
||
|
|
The call is expected to return an :class:`ROIHeads`.
|
||
|
|
"""
|
||
|
|
|
||
|
|
|
||
|
|
def build_reid_heads(cfg):
|
||
|
|
"""
|
||
|
|
Build REIDHeads defined by `cfg.MODEL.REID_HEADS.NAME`.
|
||
|
|
"""
|
||
|
|
head = cfg.MODEL.REID_HEADS.NAME
|
||
|
|
return REID_HEADS_REGISTRY.get(head)(cfg)
|