def train(args):
config.merge_from_list(args.opts)
cfg = config
model = build_model(cfg)
if not os.path.exists(cfg.OUTPUT_DIR):
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
logger.info('output will be saved into: {}'.format(cfg.OUTPUT_DIR))
trainer = Trainer(cfg, model)
trainer.resume_or_load(resume=args.resume)
if cfg.TEST.AUG.ENABLED:
trainer.register_hooks(
[hooks.EvalHook(0, lambda: trainer.test_with_TTA(cfg, trainer.model))]
)
return trainer.train()
def __init__(self, cfg):
self.cfg = cfg
self.model = build_model(self.cfg)
self.model.eval()
self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
checkpointer = DetectionCheckpointer(self.model)
print('try load weights from: {}'.format(cfg.MODEL.WEIGHTS))
checkpointer.load(cfg.MODEL.WEIGHTS)
self.transform_gen = T.ResizeShortestEdge(
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
)
self.input_format = cfg.INPUT.FORMAT
assert self.input_format in ["RGB", "BGR"], self.input_format