Spaces:
Build error
Build error
File size: 3,436 Bytes
b6af722 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
from cosmos_predict1.utils.easy_io.handlers.base import BaseFileHandler
from cosmos_predict1.utils.easy_io.handlers.csv_handler import CsvHandler
from cosmos_predict1.utils.easy_io.handlers.gzip_handler import GzipHandler
from cosmos_predict1.utils.easy_io.handlers.imageio_video_handler import ImageioVideoHandler
from cosmos_predict1.utils.easy_io.handlers.json_handler import JsonHandler
from cosmos_predict1.utils.easy_io.handlers.jsonl_handler import JsonlHandler
from cosmos_predict1.utils.easy_io.handlers.np_handler import NumpyHandler
from cosmos_predict1.utils.easy_io.handlers.pandas_handler import PandasHandler
from cosmos_predict1.utils.easy_io.handlers.pickle_handler import PickleHandler
from cosmos_predict1.utils.easy_io.handlers.pil_handler import PILHandler
from cosmos_predict1.utils.easy_io.handlers.tarfile_handler import TarHandler
from cosmos_predict1.utils.easy_io.handlers.torch_handler import TorchHandler
from cosmos_predict1.utils.easy_io.handlers.torchjit_handler import TorchJitHandler
from cosmos_predict1.utils.easy_io.handlers.txt_handler import TxtHandler
from cosmos_predict1.utils.easy_io.handlers.yaml_handler import YamlHandler
file_handlers = {
"json": JsonHandler(),
"yaml": YamlHandler(),
"yml": YamlHandler(),
"pickle": PickleHandler(),
"pkl": PickleHandler(),
"tar": TarHandler(),
"jit": TorchJitHandler(),
"npy": NumpyHandler(),
"txt": TxtHandler(),
"csv": CsvHandler(),
"pandas": PandasHandler(),
"gz": GzipHandler(),
"jsonl": JsonlHandler(),
}
for torch_type in ["pt", "pth", "ckpt"]:
file_handlers[torch_type] = TorchHandler()
for img_type in ["jpg", "jpeg", "png", "bmp", "gif"]:
file_handlers[img_type] = PILHandler()
file_handlers[img_type].format = img_type
for video_type in ["mp4", "avi", "mov", "webm", "flv", "wmv"]:
file_handlers[video_type] = ImageioVideoHandler()
def _register_handler(handler, file_formats):
"""Register a handler for some file extensions.
Args:
handler (:obj:`BaseFileHandler`): Handler to be registered.
file_formats (str or list[str]): File formats to be handled by this
handler.
"""
if not isinstance(handler, BaseFileHandler):
raise TypeError(f"handler must be a child of BaseFileHandler, not {type(handler)}")
if isinstance(file_formats, str):
file_formats = [file_formats]
if not all([isinstance(item, str) for item in file_formats]):
raise TypeError("file_formats must be a str or a list of str")
for ext in file_formats:
file_handlers[ext] = handler
def register_handler(file_formats, **kwargs):
def wrap(cls):
_register_handler(cls(**kwargs), file_formats)
return cls
return wrap
|