ICML2022/resefa

1# python3.72"""Misc utility functions."""3 4import os5import hashlib6 7from torch.hub import download_url_to_file8 9__all__ = [10    'REPO_NAME', 'Infix', 'print_and_execute', 'check_file_ext',11    'IMAGE_EXTENSIONS', 'VIDEO_EXTENSIONS', 'MEDIA_EXTENSIONS',12    'parse_file_format', 'set_cache_dir', 'get_cache_dir', 'download_url'13]14 15REPO_NAME = 'Hammer'  # Name of the repository (project).16 17 18class Infix(object):19    """Helper class to create custom infix operators.20 21    When using it, make sure to put the operator between `<<` and `>>`.22    `<< INFIX_OP_NAME >>` should be considered as a whole operator.23 24    Examples:25 26    # Use `Infix` to create infix operators directly.27    add = Infix(lambda a, b: a + b)28    1 << add >> 2  # gives 329    1 << add >> 2 << add >> 3  # gives 630 31    # Use `Infix` as a decorator.32    @Infix33    def mul(a, b):34        return a * b35    2 << mul >> 4  # gives 836    2 << mul >> 3 << mul >> 7  # gives 4237    """38 39    def __init__(self, function):40        self.function = function41        self.left_value = None42 43    def __rlshift__(self, left_value):  # override `<<` before `Infix` instance44        assert self.left_value is None  # make sure left is only called once45        self.left_value = left_value46        return self47 48    def __rshift__(self, right_value):  # override `>>` after `Infix` instance49        result = self.function(self.left_value, right_value)50        self.left_value = None  # reset to None51        return result52 53 54def print_and_execute(cmd):55    """Prints and executes a system command.56 57    Args:58        cmd: Command to be executed.59    """60    print(cmd)61    os.system(cmd)62 63 64def check_file_ext(filename, *ext_list):65    """Checks whether the given filename is with target extension(s).66 67    NOTE: If `ext_list` is empty, this function will always return `False`.68 69    Args:70        filename: Filename to check.71        *ext_list: A list of extensions.72 73    Returns:74        `True` if the filename is with one of extensions in `ext_list`,75        otherwise `False`.76    """77    if len(ext_list) == 0:78        return False79    ext_list = [ext if ext.startswith('.') else '.' + ext for ext in ext_list]80    ext_list = [ext.lower() for ext in ext_list]81    basename = os.path.basename(filename)82    ext = os.path.splitext(basename)[1].lower()83    return ext in ext_list84 85 86# File extensions regarding images (not including GIFs).87IMAGE_EXTENSIONS = (88    '.bmp', '.ppm', '.pgm', '.jpeg', '.jpg', '.jpe', '.jp2', '.png', '.webp',89    '.tiff', '.tif'90)91# File extensions regarding videos.92VIDEO_EXTENSIONS = (93    '.avi', '.mkv', '.mp4', '.m4v', '.mov', '.webm', '.flv', '.rmvb', '.rm',94    '.3gp'95)96# File extensions regarding media, i.e., images, videos, GIFs.97MEDIA_EXTENSIONS = ('.gif', *IMAGE_EXTENSIONS, *VIDEO_EXTENSIONS)98 99 100def parse_file_format(path):101    """Parses the file format of a given path.102 103    This function basically parses the file format according to its extension.104    It will also return `dir` is the given path is a directory.105 106    Parable file formats:107 108    - zip: with `.zip` extension.109    - tar: with `.tar` / `.tgz` / `.tar.gz` extension.110    - lmdb: a folder ending with `lmdb`.111    - txt: with `.txt` / `.text` extension, OR without extension (e.g. LICENSE).112    - json: with `.json` extension.113    - jpg: with `.jpeg` / `jpg` / `jpe` extension.114    - png: with `.png` extension.115 116    Args:117        path: The path to the file to parse format from.118 119    Returns:120        A lower-case string, indicating the file format, or `None` if the format121            cannot be successfully parsed.122    """123    # Handle directory.124    if os.path.isdir(path) or path.endswith('/'):125        if path.rstrip('/').lower().endswith('lmdb'):126            return 'lmdb'127        return 'dir'128    # Handle file.129    if os.path.isfile(path) and os.path.splitext(path)[1] == '':130        return 'txt'131    path = path.lower()132    if path.endswith('.tar.gz'):  # Cannot parse accurate extension.133        return 'tar'134    ext = os.path.splitext(path)[1]135    if ext == '.zip':136        return 'zip'137    if ext in ['.tar', '.tgz']:138        return 'tar'139    if ext in ['.txt', '.text']:140        return 'txt'141    if ext == '.json':142        return 'json'143    if ext in ['.jpeg', '.jpg', '.jpe']:144        return 'jpg'145    if ext == '.png':146        return 'png'147    # Unparsable.148    return None149 150 151_cache_dir = None152 153 154def set_cache_dir(directory=None):155    """Sets the global cache directory.156 157    The cache directory can be used to save some files that will be shared158    across jobs. The default cache directory is set as `~/.cache/${REPO_NAME}/`.159    This function can be used to redirect the cache directory. Or, users can use160    `None` to reset the cache directory back to default.161 162    Args:163        directory: The target directory used to cache files. If set as `None`,164            the cache directory will be reset back to default. (default: None)165    """166    assert directory is None or isinstance(directory, str), 'Invalid directory!'167    global _cache_dir  # pylint: disable=global-statement168    _cache_dir = directory169 170 171def get_cache_dir():172    """Gets the global cache directory.173 174    The global cache directory is primarily set as `~/.cache/${REPO_NAME}/` by175    default, and can be redirected with `set_cache_dir()`.176 177    Returns:178        A string, representing the global cache directory.179    """180    if _cache_dir is None:181        home = os.path.expanduser('~')182        return os.path.join(home, '.cache', REPO_NAME)183    return _cache_dir184 185 186def download_url(url, path=None, filename=None, sha256=None):187    """Downloads file from URL.188 189    This function downloads a file from given URL, and executes Hash check if190    needed.191 192    Args:193        url: The URL to download file from.194        path: Path (directory) to save the downloaded file. If set as `None`,195            the cache directory will be used. Please see `get_cache_dir()` for196            more details. (default: None)197        filename: The name to save the file. If set as `None`, this name will be198            automatically parsed from the given URL. (default: None)199        sha256: The expected sha256 of the downloaded file. If set as `None`,200            the hash check will be skipped. Otherwise, this function will check201            whether the sha256 of the downloaded file matches this field.202 203    Returns:204        A two-element tuple, where the first term is the full path of the205            downloaded file, and the second term indicate the hash check result.206            `True` means hash check passes, `False` means hash check fails,207            while `None` means no hash check is executed.208    """209    # Handle file path.210    if path is None:211        path = get_cache_dir()212    if filename is None:213        filename = os.path.basename(url)214    save_path = os.path.join(path, filename)215    # Download file if needed.216    if not os.path.exists(save_path):217        print(f'Downloading URL `{url}` to path `{save_path}` ...')218        os.makedirs(path, exist_ok=True)219        download_url_to_file(url, save_path, hash_prefix=None, progress=True)220    # Check hash if needed.221    check_result = None222    if sha256 is not None:223        with open(save_path, 'rb') as f:224            file_hash = hashlib.sha256(f.read())225            check_result = (file_hash.hexdigest() == sha256)226 227    return save_path, check_result228