Add efficientdet models (Layout-Parser#67)

* add effdet check

* Add effdet models and catalogs

* clean-up

* register effdet models

* Add generalized image type support for layout models

* Add effdet tests

* Update reqs

Author: lolipopshock

dev-requirements.txt

@@ -1,4 +1,5 @@
 pytest
+torch
 numpy
 opencv-python
 pandas
@@ -11,4 +12,5 @@ google-cloud-vision==1
 pytesseract
 pycocotools
 git+https://github.com/facebookresearch/[email protected]#egg=detectron2
-paddlepaddle
+paddlepaddle
+effdet

setup.py

@@ -24,7 +24,6 @@
         "pandas",
         "pillow",
         "pyyaml>=5.1",
-        "torch",
         "torchvision",
         "iopath",
       ],
@@ -33,6 +32,10 @@
           'google-cloud-vision==1',
           'pytesseract'
         ], 
+        "effdet": [
+          "torch",
+          "effdet"
+        ]
       },
       include_package_data=True
       )

src/layoutparser/__init__.py

@@ -6,6 +6,7 @@
     _LazyModule,
     is_detectron2_available,
     is_paddle_available,
+    is_effdet_available,
     is_pytesseract_available,
     is_gcv_available,
 )
@@ -45,6 +46,9 @@
 if is_paddle_available():
     _import_structure["models.paddledetection"] = ["PaddleDetectionLayoutModel"]
 
+if is_effdet_available():
+    _import_structure["models.effdet"] = ["EfficientDetLayoutModel"]
+
 if is_pytesseract_available():
     _import_structure["ocr.tesseract_agent"] = [
         "TesseractAgent",

src/layoutparser/file_utils.py

@@ -39,11 +39,18 @@
     # The name of the paddlepaddle library:
     # Install name: pip install paddlepaddle
     # Import name: import paddle
-    _paddle_version = importlib_metadata.version("paddlepaddle") 
+    _paddle_version = importlib_metadata.version("paddlepaddle")
     logger.debug(f"Paddle version {_paddle_version} available.")
 except importlib_metadata.PackageNotFoundError:
     _paddle_available = False
 
+_effdet_available = importlib.util.find_spec("effdet") is not None
+try:
+    _effdet_version = importlib_metadata.version("effdet")
+    logger.debug(f"Effdet version {_effdet_version} available.")
+except importlib_metadata.PackageNotFoundError:
+    _effdet_version = False
+
 ###########################################
 ############## OCR Tool Deps ##############
 ###########################################
@@ -78,12 +85,16 @@ def is_torch_cuda_available():
         return False
 
 
+def is_detectron2_available():
+    return _detectron2_available
+
+
 def is_paddle_available():
     return _paddle_available
 
 
-def is_detectron2_available():
-    return _detectron2_available
+def is_effdet_available():
+    return _effdet_available
 
 
 def is_pytesseract_available():
@@ -111,6 +122,11 @@ def is_gcv_available():
 installation page: https://github.com/PaddlePaddle/Paddle and follow the ones that match your environment.
 """
 
+EFFDET_IMPORT_ERROR = """
+{0} requires the effdet library but it was not found in your environment. You can install it with pip:
+`pip install effdet`
+"""
+
 PYTESSERACT_IMPORT_ERROR = """
 {0} requires the PyTesseract library but it was not found in your environment. You can install it with pip:
 `pip install pytesseract`
@@ -126,6 +142,7 @@ def is_gcv_available():
         ("torch", (is_torch_available, PYTORCH_IMPORT_ERROR)),
         ("detectron2", (is_detectron2_available, DETECTRON2_IMPORT_ERROR)),
         ("paddle", (is_paddle_available, PADDLE_IMPORT_ERROR)),
+        ("effdet", (is_effdet_available, )),
         ("pytesseract", (is_pytesseract_available, PYTESSERACT_IMPORT_ERROR)),
         ("google-cloud-vision", (is_gcv_available, GCV_IMPORT_ERROR)),
     ]
@@ -172,7 +189,7 @@ def __init__(
         self._import_structure = import_structure
 
         # Following [PEP 366](https://www.python.org/dev/peps/pep-0366/)
-        # The __package__ variable should be set 
+        # The __package__ variable should be set
         # https://docs.python.org/3/reference/import.html#__package__
         self.__package__ = self.__name__
 
@@ -198,4 +215,4 @@ def _get_module(self, module_name: str):
         return importlib.import_module("." + module_name, self.__name__)
 
     def __reduce__(self):
-        return (self.__class__, (self._name, self.__file__, self._import_structure))
+        return (self.__class__, (self._name, self.__file__, self._import_structure))

src/layoutparser/models/__init__.py

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

src/layoutparser/models/base_layoutmodel.py

@@ -1,17 +1,23 @@
+from typing import Union
 from abc import ABC, abstractmethod
 
 from ..file_utils import requires_backends
 
 
 class BaseLayoutModel(ABC):
-    
     @property
     @abstractmethod
     def DETECTOR_NAME(self):
         pass
-    
+
+    @abstractmethod
+    def detect(self, image):
+        pass
+
     @abstractmethod
-    def detect(self):
+    def image_loader(self, image: Union["ndarray", "Image"]):
+        """It will process the input images appropriately to the target format. 
+        """
         pass
 
     # Add lazy loading mechanisms for layout models, refer to

src/layoutparser/models/detectron2/layoutmodel.py

@@ -1,3 +1,4 @@
+from typing import Union
 from PIL import Image
 import numpy as np
 
@@ -41,7 +42,7 @@ class Detectron2LayoutModel(BaseLayoutModel):
 
     Examples::
         >>> import layoutparser as lp
-        >>> model = lp.models.Detectron2LayoutModel('lp://HJDataset/faster_rcnn_R_50_FPN_3x/config')
+        >>> model = lp.Detectron2LayoutModel('lp://HJDataset/faster_rcnn_R_50_FPN_3x/config')
         >>> model.detect(image)
 
     """
@@ -108,7 +109,7 @@ def _reconstruct_path_with_detector_name(self, path: str) -> str:
             model_name_segments = model_name.split("/")
             if (
                 len(model_name_segments) == 3
-                and "detectron2" not in model_name_segments
+                and self.DETECTOR_NAME not in model_name_segments
             ):
                 return "lp://" + self.DETECTOR_NAME + "/" + path[len("lp://") :]
         return path
@@ -148,12 +149,16 @@ def detect(self, image):
             :obj:`~layoutparser.Layout`: The detected layout of the input image
         """
 
+        image = self.image_loader(image)
+        outputs = self.model(image)
+        layout = self.gather_output(outputs)
+        return layout
+
+    def image_loader(self, image: Union["np.ndarray", "Image.Image"]):
         # Convert PIL Image Input
         if isinstance(image, Image.Image):
             if image.mode != "RGB":
                 image = image.convert("RGB")
             image = np.array(image)
 
-        outputs = self.model(image)
-        layout = self.gather_output(outputs)
-        return layout
+        return image

src/layoutparser/models/effdet/__init__.py

@@ -0,0 +1,2 @@
+from . import catalog as _UNUSED
+from .layoutmodel import EfficientDetLayoutModel

src/layoutparser/models/effdet/catalog.py

@@ -0,0 +1,53 @@
+from iopath.common.file_io import PathHandler
+
+from ..base_catalog import PathManager
+
+MODEL_CATALOG = {
+    "PubLayNet": {
+        "tf_efficientdet_d0": "https://www.dropbox.com/s/ukbw5s673633hsw/publaynet-tf_efficientdet_d0.pth.tar?dl=1",
+        "tf_efficientdet_d1": "https://www.dropbox.com/s/gxy11xkkiwnpgog/publaynet-tf_efficientdet_d1.pth.tar?dl=1"
+    },
+    "MFD": {
+        "tf_efficientdet_d0": "https://www.dropbox.com/s/dkr22iux7thlhel/mfd-tf_efficientdet_d0.pth.tar?dl=1",
+        "tf_efficientdet_d1": "https://www.dropbox.com/s/icmbiaqr5s9bz1x/mfd-tf_efficientdet_d1.pth.tar?dl=1"
+    }
+}
+
+# In effdet training scripts, it requires the label_map starting
+# from 1 instead of 0
+LABEL_MAP_CATALOG = {
+    "PubLayNet": {
+        1: "Text", 
+        2: "Title", 
+        3: "List", 
+        4: "Table", 
+        5: "Figure"
+    }
+}
+
+class LayoutParserEfficientDetModelHandler(PathHandler):
+    """
+    Resolve anything that's in LayoutParser model zoo.
+    """
+
+    PREFIX = "lp://efficientdet/"
+
+    def _get_supported_prefixes(self):
+        return [self.PREFIX]
+
+    def _get_local_path(self, path, **kwargs):
+        model_name = path[len(self.PREFIX) :]
+
+        dataset_name, *model_name, data_type = model_name.split("/")
+
+        if data_type == "weight":
+            model_url = MODEL_CATALOG[dataset_name]["/".join(model_name)]
+        else:
+            raise ValueError(f"Unknown data_type {data_type}")
+        return PathManager.get_local_path(model_url, **kwargs)
+
+    def _open(self, path, mode="r", **kwargs):
+        return PathManager.open(self._get_local_path(path), mode, **kwargs)
+
+
+PathManager.register_handler(LayoutParserEfficientDetModelHandler())