Improve the directory structure for models

Author: lolipopshock

src/layoutparser/models/__init__.py

@@ -0,0 +1,4 @@
+from . import catalog as _UNUSED
+# A trick learned from 
+# https://github.com/facebookresearch/detectron2/blob/62cf3a2b6840734d2717abdf96e2dd57ed6612a6/detectron2/checkpoint/__init__.py#L6
+from .layoutmodel import Detectron2LayoutModel

src/layoutparser/models/catalog.py

@@ -0,0 +1,57 @@
+from fvcore.common.file_io import PathHandler, PathManager, HTTPURLHandler
+
+MODEL_CATALOG = {
+    'HJDataset': {
+        'faster_rcnn_R_50_FPN_3x': 'https://www.dropbox.com/s/6icw6at8m28a2ho/model_final.pth?dl=1',
+        'mask_rcnn_R_50_FPN_3x': 'https://www.dropbox.com/s/893paxpy5suvlx9/model_final.pth?dl=1',
+        'retinanet_R_50_FPN_3x': 'https://www.dropbox.com/s/yxsloxu3djt456i/model_final.pth?dl=1'
+    }
+}
+
+CONFIG_CATALOG = {
+    'HJDataset': {
+        'faster_rcnn_R_50_FPN_3x': 'https://www.dropbox.com/s/j4yseny2u0hn22r/config.yml?dl=1',
+        'mask_rcnn_R_50_FPN_3x': 'https://www.dropbox.com/s/4jmr3xanmxmjcf8/config.yml?dl=0',
+        'retinanet_R_50_FPN_3x': 'https://www.dropbox.com/s/z8a8ywozuyc5c2x/config.yml?dl=0'
+    }
+}
+
+class DropboxHandler(HTTPURLHandler):
+    """
+    Supports download and file check for dropbox links
+    """
+    def _get_supported_prefixes(self):
+        return ["https://www.dropbox.com"]
+
+    def _isfile(self, path):
+        return path in self.cache_map
+
+
+class LayoutParserHandler(PathHandler):
+    """
+    Resolve anything that's in LayoutParser model zoo.
+    """
+    
+    PREFIX = "lp://"
+
+    def _get_supported_prefixes(self):
+        return [self.PREFIX]
+
+    def _get_local_path(self, path):
+        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)]
+        elif data_type == 'config':
+            model_url = CONFIG_CATALOG[dataset_name]['/'.join(model_name)]
+        else:
+            raise ValueError(f"Unknown data_type {data_type}")
+        return PathManager.get_local_path(model_url)
+
+    def _open(self, path, mode="r", **kwargs):
+        return PathManager.open(self._get_local_path(path), mode, **kwargs)
+
+
+PathManager.register_handler(DropboxHandler())
+PathManager.register_handler(LayoutParserHandler())

src/layoutparser/models.py src/layoutparser/models/layoutmodel.pysrc/layoutparser/models.py renamed to src/layoutparser/models/layoutmodel.py

@@ -3,7 +3,9 @@
 import torch
 from detectron2.config import get_cfg
 from detectron2.engine import DefaultPredictor
-from .elements import *
+from ..elements import *
+from fvcore.common.file_io import PathManager
+
 
 __all__ = ['Detectron2LayoutModel']
 
@@ -16,13 +18,14 @@ def detect(self): pass
 
 class Detectron2LayoutModel(BaseLayoutModel):
 
-    def __init__(self, config_name,
+    def __init__(self, config_path,
                        model_path = None,
                        label_map  = None,
                        extra_config= []):
 
         cfg = get_cfg()
-        cfg.merge_from_file(config_name)
+        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: