Source code for hfai.utils.detr2.ffrecord_handler

import os, io, sys
os.environ["DETECTRON2_DATASETS"] = "ffrecord://"

from iopath.common.file_io import PathHandler
from detectron2.utils.file_io import PathManager
from ffrecord import PackedFolder


class FFRecordHandler(PathHandler):

    PREFIX = "ffrecord://"

    def __init__(self, ffr_file, ffr_prefix, local_prefix):
        super().__init__()
        self.ffr_file = ffr_file
        self.folder = PackedFolder(self.ffr_file)
        self.ffr_prefix = ffr_prefix
        self.local_prefix = local_prefix

    def _get_supported_prefixes(self):
        return [self.PREFIX]

    def _get_local_path(self, path, **kwargs):
        name = path[len(self.PREFIX):]
        if self.local_prefix:
            name = os.path.join(self.local_prefix, name)
        return name

    def _open(self, path, mode="r", **kwargs):
        assert self.folder is not None
        path = path[len(self.PREFIX):]
        if self.ffr_prefix:
            path = os.path.relpath(path, self.ffr_prefix)
        data = self.folder.read_one(path)

        if mode == "rb":
            return io.BytesIO(data)
        elif mode == "r":
            encoding = kwargs.pop("encoding", sys.getfilesystemencoding())
            text = str(data, encoding=encoding)
            return io.StringIO(text)
        else:
            raise ValueError(f"unsupported mode: {mode}")


def register_ffrecord_handler(ffr_file, ffr_prefix=None, local_prefix="datasets"):
    """
    在 detectron2 上使用 FFRecord，需要在导入 detectron2 包之前使用本函数

    具体教程请看 `HFAI X Detectron2 <https://doc.hfai.high-flyer.cn/performance/detr2.html>`_

    Examples:

    .. code-block:: python

        from hfai.utils.detr2 import register_ffrecord_handler
        register_ffrecord_handler(
            ffr_file="datasets/coco/coco.ffr",
            ffr_prefix="coco",
        )

        # 在这之后我们导入 detectron2 并使用
        import detectron

    """
    ffr_handler = FFRecordHandler(ffr_file, ffr_prefix, local_prefix)
    PathManager.register_handler(ffr_handler)