Actually make multi-model segmentation work

mittagessen · Jul 29, 2024 · 38484c3 · 38484c3
1 parent 4219c46
commit 38484c3
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/kraken/kraken.py b/kraken/kraken.py
@@ -501,8 +501,6 @@ def segment(ctx, model, boxes, text_direction, scale, maxcolseps,
     """
     from kraken.containers import ProcessingStep
 
-    print(model)
-
     if model and boxes:
         logger.warning(f'Baseline model ({model}) given but legacy segmenter selected. Forcing to -bl.')
         boxes = False
@@ -536,7 +534,8 @@ def segment(ctx, model, boxes, text_direction, scale, maxcolseps,
         for loc in locations:
             message(f'Loading ANN {loc}\t', nl=False)
             try:
-                models.append(TorchVGSLModel.load_model(loc).to(ctx.meta['device']))
+                models.append(TorchVGSLModel.load_model(loc))
+                models[-1].to(ctx.meta['device'])
             except Exception:
                 if ctx.meta['raise_failed']:
                     raise