[feat] AutoLayoutModel and flexible model configs (#69)

* move _reconstruct_path_with_detector_name to the baseclass

* improve the installation logic

* allow three model config formats

* update tests

* minor fixes

* Add complete tests

* Add AutoLayoutModel

* Better error message

* Move to automodel folder

* warnings for enforce_cpu

* Add AutoModel tests

* cleanup

Author: lolipopshock

setup.py

@@ -31,17 +31,29 @@
           'google-cloud-vision==1',
           'pytesseract'
         ], 
+        "gcv": [
+          'google-cloud-vision==1',
+        ],
+        "tesseract": [
+          'pytesseract'
+        ],
+        "layoutmodels": [
+          "torch",
+          "torchvision",
+          "effdet"
+        ],
         "effdet": [
           "torch",
           "torchvision",
           "effdet"
         ],
         "detectron2": [
-          "detectron2@git+https://github.com/facebookresearch/[email protected]#egg=detectron2"
+          "detectron2@git+https://github.com/facebookresearch/[email protected]#egg=detectron2"
+          # Supporting detectron0.5 for compatibility with newer torch versions 
         ],
         "paddledetection": [
           "paddlepaddle==2.1.0"
-        ]
+        ],
       },
       include_package_data=True
       )

src/layoutparser/__init__.py

@@ -40,6 +40,8 @@
     ]
 }
 
+_import_structure["models"] = ["AutoLayoutModel"]
+
 if is_detectron2_available():
     _import_structure["models.detectron2"] = ["Detectron2LayoutModel"]
 

src/layoutparser/models/__init__.py

@@ -1,3 +1,4 @@
 from .detectron2.layoutmodel import Detectron2LayoutModel
 from .paddledetection.layoutmodel import PaddleDetectionLayoutModel
-from .effdet.layoutmodel import EfficientDetLayoutModel
+from .effdet.layoutmodel import EfficientDetLayoutModel
+from .auto_layoutmodel import AutoLayoutModel

src/layoutparser/models/auto_layoutmodel.py

@@ -0,0 +1,56 @@
+from typing import Optional, Dict, Union, List
+from .detectron2.layoutmodel import Detectron2LayoutModel
+from .paddledetection.layoutmodel import PaddleDetectionLayoutModel
+from .effdet.layoutmodel import EfficientDetLayoutModel
+from .model_config import (
+    is_lp_layout_model_config_any_format,
+)
+
+ALL_AVAILABLE_BACKENDS = {
+    Detectron2LayoutModel.DETECTOR_NAME: Detectron2LayoutModel,
+    PaddleDetectionLayoutModel.DETECTOR_NAME: PaddleDetectionLayoutModel,
+    EfficientDetLayoutModel.DETECTOR_NAME: EfficientDetLayoutModel,
+}
+
+
+def AutoLayoutModel(
+    config_path: str,
+    model_path: Optional[str] = None,
+    label_map: Optional[Dict]=None,
+    device: Optional[str]=None,
+    extra_config: Optional[Union[Dict, List]]=None,
+) -> "BaseLayoutModel":
+    """[summary]
+
+    Args:
+        config_path (:obj:`str`):
+            The path to the configuration file.
+        model_path (:obj:`str`, None):
+            The path to the saved weights of the model.
+            If set, overwrite the weights in the configuration file.
+            Defaults to `None`.
+        label_map (:obj:`dict`, optional):
+            The map from the model prediction (ids) to real
+            word labels (strings). If the config is from one of the supported
+            datasets, Layout Parser will automatically initialize the label_map.
+            Defaults to `None`.
+        device(:obj:`str`, optional):
+            Whether to use cuda or cpu devices. If not set, LayoutParser will
+            automatically determine the device to initialize the models on. 
+        extra_config (:obj:`dict`, optional):
+            Extra configuration passed used for initializing the layout model.
+
+    Returns:
+        # BaseLayoutModel: the create LayoutModel instance
+    """
+    if not is_lp_layout_model_config_any_format(config_path):
+        raise ValueError(f"Invalid model config_path {config_path}")
+    for backend_name in ALL_AVAILABLE_BACKENDS:
+        if backend_name in config_path:
+            return ALL_AVAILABLE_BACKENDS[backend_name](
+                config_path,
+                model_path=model_path,
+                label_map=label_map,
+                extra_config=extra_config,
+                device=device,
+            )

src/layoutparser/models/base_layoutmodel.py

@@ -1,33 +1,72 @@
-from typing import Union
+from typing import Optional, Tuple, Union, Dict
 from abc import ABC, abstractmethod
 
+from .model_config import LayoutModelConfig, add_identifier_for_config, layout_model_config_parser, is_lp_layout_model_config_any_format
 from ..file_utils import requires_backends
 
-
 class BaseLayoutModel(ABC):
+
+    # TODO: Build a metaclass for lazy module loader
+    @property
+    @abstractmethod
+    def DEPENDENCIES(self):
+        """DEPENDENCIES lists all necessary dependencies for the class."""
+        pass
+
     @property
     @abstractmethod
     def DETECTOR_NAME(self):
         pass
 
+    @property
     @abstractmethod
-    def detect(self, image):
+    def MODEL_CATALOG(self) -> Dict[str, Dict[str, str]]:
         pass
 
     @abstractmethod
-    def image_loader(self, image: Union["ndarray", "Image"]):
-        """It will process the input images appropriately to the target format. 
-        """
+    def detect(self, image: Union["np.ndarray", "Image.Image"]):
         pass
 
-    # Add lazy loading mechanisms for layout models, refer to
-    # layoutparser.ocr.BaseOCRAgent
-    # TODO: Build a metaclass for lazy module loader
-    @property
+
     @abstractmethod
-    def DEPENDENCIES(self):
-        """DEPENDENCIES lists all necessary dependencies for the class."""
+    def image_loader(self, image: Union["np.ndarray", "Image.Image"]):
+        """It will process the input images appropriately to the target format."""
         pass
+    
+    def _parse_config(self, config_path:str, identifier:str) -> Union[LayoutModelConfig, str]:
+        
+        if is_lp_layout_model_config_any_format(config_path):
+            config_path = add_identifier_for_config(config_path, identifier)
+            for dataset_name in self.MODEL_CATALOG:
+                if dataset_name in config_path:
+                    default_model_arch = list(self.MODEL_CATALOG[dataset_name].keys())[0]
+                    # Use the first model_name for the dataset as the default_model_arch
+                    return layout_model_config_parser(config_path, self.DETECTOR_NAME, default_model_arch)
+            raise ValueError(f"The config {config_path} is not a valid config for {self.__class__}, "
+                             f"possibly because there aren't models trained for the specified dataset.")
+        else:
+            return config_path
+
+    def config_parser(self, config_path:str, model_path: Optional[str], allow_empty_path=False) -> Tuple[str, str]:
+
+        config_path = self._parse_config(config_path, "config")
+        
+        if isinstance(config_path, str) and model_path is None:
+            if not allow_empty_path:
+                raise ValueError(
+                    f"Invalid config and model path pairs ({(config_path, model_path)}):"
+                    f"When config_path is a regular URL, the model_path should not be empty"
+                )
+            else:
+                return config_path, model_path
+        elif isinstance(config_path, LayoutModelConfig) and model_path is None:
+            model_path = config_path.dual()
+        else:
+            model_path = self._parse_config(model_path, "weight")
+
+        config_path = config_path if isinstance(config_path, str) else config_path.full
+        model_path = model_path if isinstance(model_path, str) else model_path.full
+        return config_path, model_path
 
     def __new__(cls, *args, **kwargs):
 

src/layoutparser/models/detectron2/layoutmodel.py

@@ -1,8 +1,9 @@
 from typing import Union
 from PIL import Image
 import numpy as np
+import warnings
 
-from .catalog import PathManager, LABEL_MAP_CATALOG
+from .catalog import MODEL_CATALOG, PathManager, LABEL_MAP_CATALOG
 from ..base_layoutmodel import BaseLayoutModel
 from ...elements import Rectangle, TextBlock, Layout
 from ...file_utils import is_torch_cuda_available, is_detectron2_available
@@ -30,9 +31,9 @@ class Detectron2LayoutModel(BaseLayoutModel):
             word labels (strings). If the config is from one of the supported
             datasets, Layout Parser will automatically initialize the label_map.
             Defaults to `None`.
-        enforce_cpu(:obj:`bool`, optional):
-            When set to `True`, it will enforce using cpu even if it is on a CUDA
-            available device.
+        device(:obj:`str`, optional):
+            Whether to use cuda or cpu devices. If not set, LayoutParser will
+            automatically determine the device to initialize the models on.
         extra_config (:obj:`list`, optional):
             Extra configuration passed to the Detectron2 model
             configuration. The argument will be used in the `merge_from_list
@@ -49,70 +50,55 @@ class Detectron2LayoutModel(BaseLayoutModel):
 
     DEPENDENCIES = ["detectron2"]
     DETECTOR_NAME = "detectron2"
+    MODEL_CATALOG = MODEL_CATALOG
 
     def __init__(
         self,
         config_path,
         model_path=None,
         label_map=None,
         extra_config=None,
-        enforce_cpu=False,
+        enforce_cpu=None,
+        device=None,
     ):
 
+        if enforce_cpu is not None:
+            warnings.warn(
+                "Setting enforce_cpu is deprecated. Please set `device` instead.",
+                DeprecationWarning,
+            )
+
         if extra_config is None:
             extra_config = []
 
-        if config_path.startswith("lp://") and label_map is None:
-            dataset_name = config_path.lstrip("lp://").split("/")[0]
-            label_map = LABEL_MAP_CATALOG[dataset_name]
-
-        if enforce_cpu:
-            extra_config.extend(["MODEL.DEVICE", "cpu"])
+        config_path, model_path = self.config_parser(
+            config_path, model_path, allow_empty_path=True
+        )
+        config_path = PathManager.get_local_path(config_path)
 
         cfg = detectron2.config.get_cfg()
-        config_path = self._reconstruct_path_with_detector_name(config_path)
-        config_path = PathManager.get_local_path(config_path)
         cfg.merge_from_file(config_path)
         cfg.merge_from_list(extra_config)
 
         if model_path is not None:
-            model_path = self._reconstruct_path_with_detector_name(model_path)
+            model_path = PathManager.get_local_path(model_path)
+            # Because it will be forwarded to the detectron2 paths
             cfg.MODEL.WEIGHTS = model_path
-        
-        if not enforce_cpu:
-            cfg.MODEL.DEVICE = "cuda" if is_torch_cuda_available() else "cpu"
+
+        if is_torch_cuda_available():
+            if device is None:
+                device = "cuda"
+        else:
+            device = "cpu"
+        cfg.MODEL.DEVICE = device
 
         self.cfg = cfg
 
         self.label_map = label_map
         self._create_model()
 
-    def _reconstruct_path_with_detector_name(self, path: str) -> str:
-        """This function will add the detector name (detectron2) into the
-        lp model config path to get the "canonical" model name.
-
-        For example, for a given config_path `lp://HJDataset/faster_rcnn_R_50_FPN_3x/config`,
-        it will transform it into `lp://detectron2/HJDataset/faster_rcnn_R_50_FPN_3x/config`.
-        However, if the config_path already contains the detector name, we won't change it.
-
-        This function is a general step to support multiple backends in the layout-parser
-        library.
-
-        Args:
-            path (str): The given input path that might or might not contain the detector name.
-
-        Returns:
-            str: a modified path that contains the detector name.
-        """
-        if path.startswith("lp://"):  # TODO: Move "lp://" to a constant
-            model_name = path[len("lp://") :]
-            model_name_segments = model_name.split("/")
-            if (
-                len(model_name_segments) == 3
-                and self.DETECTOR_NAME not in model_name_segments
-            ):
-                return "lp://" + self.DETECTOR_NAME + "/" + path[len("lp://") :]
-        return path
+    def _create_model(self):
+        self.model = detectron2.engine.DefaultPredictor(self.cfg)
 
     def gather_output(self, outputs):
 
@@ -136,9 +122,6 @@ def gather_output(self, outputs):
 
         return layout
 
-    def _create_model(self):
-        self.model = detectron2.engine.DefaultPredictor(self.cfg)
-
     def detect(self, image):
         """Detect the layout of a given image.
 

src/layoutparser/models/effdet/catalog.py

@@ -22,6 +22,9 @@
         3: "List", 
         4: "Table", 
         5: "Figure"
+    },
+    "MFD": {
+        1: "Equation",
     }
 }