未验证 提交 4f6ffb40 编写于 作者: G Guanghua Yu 提交者: GitHub

fix use vdl bug in infer and export model (#2365)

上级 7f4a1dea
...@@ -101,7 +101,7 @@ class Trainer(object): ...@@ -101,7 +101,7 @@ class Trainer(object):
def _init_callbacks(self): def _init_callbacks(self):
if self.mode == 'train': if self.mode == 'train':
self._callbacks = [LogPrinter(self), Checkpointer(self)] self._callbacks = [LogPrinter(self), Checkpointer(self)]
if self.cfg.use_vdl: if 'use_vdl' in self.cfg and self.cfg.use_vdl:
self._callbacks.append(VisualDLWriter(self)) self._callbacks.append(VisualDLWriter(self))
self._compose_callback = ComposeCallback(self._callbacks) self._compose_callback = ComposeCallback(self._callbacks)
elif self.mode == 'eval': elif self.mode == 'eval':
...@@ -109,7 +109,7 @@ class Trainer(object): ...@@ -109,7 +109,7 @@ class Trainer(object):
if self.cfg.metric == 'WiderFace': if self.cfg.metric == 'WiderFace':
self._callbacks.append(WiferFaceEval(self)) self._callbacks.append(WiferFaceEval(self))
self._compose_callback = ComposeCallback(self._callbacks) self._compose_callback = ComposeCallback(self._callbacks)
elif self.mode == 'test' and self.cfg.use_vdl: elif self.mode == 'test' and 'use_vdl' in self.cfg and self.cfg.use_vdl:
self._callbacks = [VisualDLWriter(self)] self._callbacks = [VisualDLWriter(self)]
self._compose_callback = ComposeCallback(self._callbacks) self._compose_callback = ComposeCallback(self._callbacks)
else: else:
...@@ -349,6 +349,7 @@ class Trainer(object): ...@@ -349,6 +349,7 @@ class Trainer(object):
int(outs['im_id']), catid2name, int(outs['im_id']), catid2name,
draw_threshold) draw_threshold)
self.status['result_image'] = np.array(image.copy()) self.status['result_image'] = np.array(image.copy())
if self._compose_callback:
self._compose_callback.on_step_end(self.status) self._compose_callback.on_step_end(self.status)
# save image with detection # save image with detection
save_name = self._get_save_image_name(output_dir, image_path) save_name = self._get_save_image_name(output_dir, image_path)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册