build.py 565 B

123456789101112131415161718
  1. from .fcos import FcosPSS
  2. from .criterion import SetCriterion
  3. def build_fcos_pss(cfg, is_val=False):
  4. # ------------ build object detector ------------
  5. ## RT-FCOS
  6. model = FcosPSS(cfg = cfg,
  7. conf_thresh = cfg.train_conf_thresh if is_val else cfg.test_conf_thresh,
  8. topk_results = cfg.train_topk if is_val else cfg.test_topk,
  9. )
  10. criterion = None
  11. if is_val:
  12. # build criterion for training
  13. criterion = SetCriterion(cfg)
  14. return model, criterion