__init__.py 678 B

1234567891011121314151617181920
  1. from .elandarknet.build import build_elandarknet
  2. from .cspdarknet.build import build_cspdarknet
  3. from .darknet.build import build_darknet
  4. from .gelan.build import build_gelan
  5. def build_model(args):
  6. # --------------------------- ResNet series ---------------------------
  7. if 'elandarknet' in args.model:
  8. model = build_elandarknet(args)
  9. elif 'cspdarknet' in args.model:
  10. model = build_cspdarknet(args)
  11. elif 'darknet' in args.model:
  12. model = build_darknet(args)
  13. elif 'gelan' in args.model:
  14. model = build_gelan(args)
  15. else:
  16. raise NotImplementedError("Unknown model: {}".format(args.model))
  17. return model