From 94a2f9fbb72b369d2fabd1957b6aa677161035fd Mon Sep 17 00:00:00 2001
From: Guanghua Yu <742925032@qq.com>
Date: Tue, 12 May 2020 21:17:43 +0800
Subject: [PATCH] fix `with_background` in export model (#633)

* fix with_background in export model

* fix coco_eval background
---
 deploy/python/infer.py   | 2 --
 ppdet/utils/coco_eval.py | 6 +++++-
 tools/export_model.py    | 1 +
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/deploy/python/infer.py b/deploy/python/infer.py
index 45d2d9051..d042d3b18 100644
--- a/deploy/python/infer.py
+++ b/deploy/python/infer.py
@@ -283,8 +283,6 @@ class Config():
         self.use_python_inference = yml_conf['use_python_inference']
         self.min_subgraph_size = yml_conf['min_subgraph_size']
         self.labels = yml_conf['label_list']
-        if not yml_conf['with_background']:
-            self.labels = self.labels[1:]
         self.mask_resolution = None
         if 'mask_resolution' in yml_conf:
             self.mask_resolution = yml_conf['mask_resolution']
diff --git a/ppdet/utils/coco_eval.py b/ppdet/utils/coco_eval.py
index 4c5a6ab08..8fde61de6 100644
--- a/ppdet/utils/coco_eval.py
+++ b/ppdet/utils/coco_eval.py
@@ -425,7 +425,9 @@ def get_category_info_from_anno(anno_file, with_background=True):
         for i, cat in enumerate(cats)
     }
     catid2name = {cat['id']: cat['name'] for cat in cats}
-
+    if with_background:
+        clsid2catid.update({0: 0})
+        catid2name.update({0: 'background'})
     return clsid2catid, catid2name
 
 
@@ -607,5 +609,7 @@ def coco17_category_info(with_background=True):
 
     if not with_background:
         clsid2catid = {k - 1: v for k, v in clsid2catid.items()}
+    else:
+        clsid2catid.update({0: 0})
 
     return clsid2catid, catid2name
diff --git a/tools/export_model.py b/tools/export_model.py
index 20b43e8c1..8cadb912d 100644
--- a/tools/export_model.py
+++ b/tools/export_model.py
@@ -61,6 +61,7 @@ def parse_reader(reader_cfg, metric, arch):
                 metric))
     clsid2catid, catid2name = get_category_info(anno_file, with_background,
                                                 use_default_label)
+
     label_list = [str(cat) for cat in catid2name.values()]
 
     sample_transforms = reader_cfg['sample_transforms']
-- 
GitLab