From 0ce7a826d44055fe773b4473d456d7cabe3a15e9 Mon Sep 17 00:00:00 2001 From: lbrgt <43443018+lbrgt@users.noreply.github.com> Date: Fri, 9 Apr 2021 11:47:24 +0200 Subject: [PATCH 1/2] Update setup.py --- setup.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 45d6129..5884c7d 100644 --- a/setup.py +++ b/setup.py @@ -27,7 +27,7 @@ "torch", "torchvision", "pycocotools>=2.0.2", - "fvcore==0.1.1.post20200623", + "fvcore>=0.1.1", ], extras_require={ "ocr": [ @@ -36,4 +36,4 @@ ], }, include_package_data=True - ) \ No newline at end of file + ) From 653701a966b3c9597ca3b134424e6998807d1c8a Mon Sep 17 00:00:00 2001 From: lucas Date: Fri, 9 Apr 2021 15:31:16 +0200 Subject: [PATCH 2/2] Add manual cuda flag --- src/layoutparser/models/layoutmodel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/layoutparser/models/layoutmodel.py b/src/layoutparser/models/layoutmodel.py index 78b6973..acbb30a 100644 --- a/src/layoutparser/models/layoutmodel.py +++ b/src/layoutparser/models/layoutmodel.py @@ -92,7 +92,7 @@ class Detectron2LayoutModel(BaseLayoutModel): {"import_name": "_config", "module_path": "detectron2.config"}, ] - def __init__(self, config_path, model_path=None, label_map=None, extra_config=[]): + def __init__(self, config_path, model_path=None, label_map=None, extra_config=[], cuda_flag=False): cfg = self._config.get_cfg() config_path = PathManager.get_local_path(config_path) @@ -101,7 +101,7 @@ def __init__(self, config_path, model_path=None, label_map=None, extra_config=[] if model_path is not None: cfg.MODEL.WEIGHTS = model_path - cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + cfg.MODEL.DEVICE = "cuda" if cuda_flag else "cpu" self.cfg = cfg self.label_map = label_map