diff --git a/kraken/kraken.py b/kraken/kraken.py index cb89b8832..47bc913ab 100644 --- a/kraken/kraken.py +++ b/kraken/kraken.py @@ -532,12 +532,12 @@ def segment(ctx, model, boxes, text_direction, scale, maxcolseps, from kraken.lib.vgsl import TorchVGSLModel - models = [] + model = [] for loc in locations: message(f'Loading ANN {loc}\t', nl=False) try: - models.append(TorchVGSLModel.load_model(loc)) - models[-1].to(ctx.meta['device']) + model.append(TorchVGSLModel.load_model(loc)) + model[-1].to(ctx.meta['device']) except Exception: if ctx.meta['raise_failed']: raise @@ -555,7 +555,7 @@ def segment(ctx, model, boxes, text_direction, scale, maxcolseps, 'remove_hlines': remove_hlines, 'pad': pad})) - return partial(segmenter, boxes, models, text_direction, scale, maxcolseps, + return partial(segmenter, boxes, model, text_direction, scale, maxcolseps, black_colseps, remove_hlines, pad, mask, ctx.meta['device']) diff --git a/tests/test_cli.py b/tests/test_cli.py index 58dc61927..4de6246ed 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -57,11 +57,3 @@ def test_segment_color(self): with tempfile.NamedTemporaryFile() as fp: result = self.runner.invoke(cli, ['-r', '-i', self.color_img, fp.name, 'segment']) self.assertEqual(result.exit_code, 1) - - def test_segment_color_noraise(self): - """ - Tests that segmentation does not return 1 when given color image in noraise mode. - """ - with tempfile.NamedTemporaryFile() as fp: - result = self.runner.invoke(cli, ['-i', self.color_img, fp.name, 'segment']) - self.assertEqual(result.exit_code, 0)