Skip to content

Commit 52ce56b

Browse files
authored
[fix] set label_map in Detectron2LayoutModel (Layout-Parser#75)
* remove unused type annotations * automatically set label_map in Detectron2LayoutModel
1 parent 73e3015 commit 52ce56b

File tree

2 files changed

+9
-3
lines changed

2 files changed

+9
-3
lines changed

src/layoutparser/models/detectron2/layoutmodel.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,13 @@ def __init__(
9090
)
9191
config_path = PathManager.get_local_path(config_path)
9292

93+
if label_map is None:
94+
if config_path.startswith("lp://"):
95+
dataset_name = config_path.lstrip("lp://").split("/")[1]
96+
label_map = LABEL_MAP_CATALOG[dataset_name]
97+
else:
98+
label_map = {}
99+
93100
cfg = detectron2.config.get_cfg()
94101
cfg.merge_from_file(config_path)
95102
cfg.merge_from_list(extra_config)
@@ -126,8 +133,8 @@ def gather_output(self, outputs):
126133
for score, box, label in zip(scores, boxes, labels):
127134
x_1, y_1, x_2, y_2 = box
128135

129-
if self.label_map is not None:
130-
label = self.label_map.get(label, label)
136+
137+
label = self.label_map.get(label, label)
131138

132139
cur_block = TextBlock(
133140
Rectangle(x_1, y_1, x_2, y_2), type=label, score=score

src/layoutparser/models/model_config.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,6 @@
3232
- Brief format: lp://<dataset-name>/<config, weight>
3333
"""
3434

35-
from typing import List, OrderedDict, Union, Dict, Any, Tuple, Optional, NamedTuple
3635
from dataclasses import dataclass
3736

3837
LAYOUT_PARSER_MODEL_PREFIX = "lp://"

0 commit comments

Comments
 (0)