catalog.py

# Copyright 2021 The Layout Parser team and Paddle Detection model 
# contributors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import logging
from typing import Any, Optional
from urllib.parse import urlparse
import tarfile
import uuid

from iopath.common.file_io import PathHandler
from iopath.common.file_io import HTTPURLHandler
from iopath.common.file_io import get_cache_dir, file_lock
from iopath.common.download import download

from ..base_catalog import PathManager

MODEL_CATALOG = {
    "PubLayNet": {
        "ppyolov2_r50vd_dcn_365e": "https://paddle-model-ecology.bj.bcebos.com/model/layout-parser/ppyolov2_r50vd_dcn_365e_publaynet.tar",
    },
    "TableBank": {
        "ppyolov2_r50vd_dcn_365e": "https://paddle-model-ecology.bj.bcebos.com/model/layout-parser/ppyolov2_r50vd_dcn_365e_tableBank_word.tar",
        # "ppyolov2_r50vd_dcn_365e_tableBank_latex": "https://paddle-model-ecology.bj.bcebos.com/model/layout-parser/ppyolov2_r50vd_dcn_365e_tableBank_latex.tar",
        # TODO: Train a single tablebank model for paddlepaddle
    },
}

# fmt: off
LABEL_MAP_CATALOG = {
    "PubLayNet": {
        0: "Text",
        1: "Title",
        2: "List",
        3: "Table",
        4: "Figure"},
    "TableBank": {
        0: "Table"
    },
}
# fmt: on


# Paddle model package everything in tar files, and each model's tar file should contain
# the following files in the list:
_TAR_FILE_NAME_LIST = [
    "inference.pdiparams",
    "inference.pdiparams.info",
    "inference.pdmodel",
]


def _get_untar_directory(tar_file: str) -> str:

    base_path = os.path.dirname(tar_file)
    file_name = os.path.splitext(os.path.basename(tar_file))[0]
    target_folder = os.path.join(base_path, file_name)

    return target_folder


def _untar_model_weights(model_tar):
    """untar model files"""

    model_dir = _get_untar_directory(model_tar)

    if not os.path.exists(
        os.path.join(model_dir, _TAR_FILE_NAME_LIST[0])
    ) or not os.path.exists(os.path.join(model_dir, _TAR_FILE_NAME_LIST[2])):
        # the path to save the decompressed file
        os.makedirs(model_dir, exist_ok=True)
        with tarfile.open(model_tar, "r") as tarobj:
            for member in tarobj.getmembers():
                filename = None
                for tar_file_name in _TAR_FILE_NAME_LIST:
                    if tar_file_name in member.name:
                        filename = tar_file_name
                if filename is None:
                    continue
                file = tarobj.extractfile(member)
                with open(os.path.join(model_dir, filename), "wb") as model_file:
                    model_file.write(file.read())
    return model_dir


def is_cached_folder_exists_and_valid(cached):
    possible_extracted_model_folder = _get_untar_directory(cached)
    if not os.path.exists(possible_extracted_model_folder):
        return False
    for tar_file in _TAR_FILE_NAME_LIST:
        if not os.path.exists(os.path.join(possible_extracted_model_folder, tar_file)):
            return False
    return True


class PaddleModelURLHandler(HTTPURLHandler):
    """
    Supports download and file check for Baidu Cloud links
    """

    MAX_FILENAME_LEN = 250

    def _get_supported_prefixes(self):
        return ["https://paddle-model-ecology.bj.bcebos.com"]

    def _isfile(self, path):
        return path in self.cache_map

    def _get_local_path(
        self,
        path: str,
        force: bool = False,
        cache_dir: Optional[str] = None,
        **kwargs: Any,
    ) -> str:
        """
        As paddle model stores all files in tar files, we need to extract them
        and get the newly extracted folder path. This function rewrites the base
        function to support the following situations:

        1. If the tar file is not downloaded, it will download the tar file,
            extract it to the target folder, delete the downloaded tar file,
            and return the folder path.
        2. If the extracted target folder is present, and all the necessary model
            files are present (specified in _TAR_FILE_NAME_LIST), it will
            return the folder path.
        3. If the tar file is downloaded, but the extracted target folder is not
            present (or it doesn't contain the necessary files in _TAR_FILE_NAME_LIST),
            it will extract the tar file to the target folder, delete the tar file,
            and return the folder path.

        """
        self._check_kwargs(kwargs)
        if (
            force
            or path not in self.cache_map
            or not os.path.exists(self.cache_map[path])
        ):
            logger = logging.getLogger(__name__)
            parsed_url = urlparse(path)
            dirname = os.path.join(
                get_cache_dir(cache_dir), os.path.dirname(parsed_url.path.lstrip("/"))
            )
            filename = path.split("/")[-1]
            if len(filename) > self.MAX_FILENAME_LEN:
                filename = filename[:100] + "_" + uuid.uuid4().hex

            cached = os.path.join(dirname, filename)

            if is_cached_folder_exists_and_valid(cached):
                # When the cached folder exists and valid, we don't need to redownload
                # the tar file.
                self.cache_map[path] = _get_untar_directory(cached)

            else:
                with file_lock(cached):
                    if not os.path.isfile(cached):
                        logger.info("Downloading {} ...".format(path))
                        cached = download(path, dirname, filename=filename)

                    if path.endswith(".tar"):
                        model_dir = _untar_model_weights(cached)
                        try:
                            os.remove(cached)  # remove the redundant tar file
                            # TODO: remove the .lock file .
                        except:
                            logger.warning(
                                f"Not able to remove the cached tar file {cached}"
                            )

                logger.info("URL {} cached in {}".format(path, model_dir))
                self.cache_map[path] = model_dir

        return self.cache_map[path]


class LayoutParserPaddleModelHandler(PathHandler):
    """
    Resolve anything that's in LayoutParser model zoo.
    """

    PREFIX = "lp://paddledetection/"

    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(PaddleModelURLHandler())
PathManager.register_handler(LayoutParserPaddleModelHandler())