from .yolof_head     import YolofHead
from .fcos_head      import FcosHead


# build head
def build_head(cfg, in_dim, out_dim, num_classes):
    print('==============================')
    print('Head: {}'.format(cfg['head']))
    
    if cfg['head'] == 'fcos_head':
        model = FcosHead(cfg          = cfg,
                         in_dim       = in_dim,
                         out_dim      = out_dim,
                         num_classes  = num_classes,
                         num_cls_head = cfg['num_cls_head'],
                         num_reg_head = cfg['num_reg_head'],
                         act_type     = cfg['head_act'],
                         norm_type    = cfg['head_norm']
                         )
    elif cfg['head'] == 'yolof_head':
        model = YolofHead(cfg          = cfg,
                          in_dim       = in_dim,
                          out_dim      = out_dim,
                          num_classes  = num_classes,
                          num_cls_head = cfg['num_cls_head'],
                          num_reg_head = cfg['num_reg_head'],
                          act_type     = cfg['head_act'],
                          norm_type    = cfg['head_norm']
                          )

    return model