utils/io_utils.py

# -*- coding: utf-8 -*-

import os
import errno
import pickle
from collections import deque
import torch


def mkdirs(path):
    try:
        os.makedirs(path)
    except OSError as exc:  # Python >2.5
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise


def pickle_load(file_path):
    assert os.path.isfile(file_path), file_path
    with open(file_path, 'rb') as f:
        file = pickle.load(f)
    return file


def pickle_dump(file_path, file):
    if os.path.isfile(file_path):
        os.remove(file_path)
    with open(file_path, 'wb') as f:
        pickle.dump(file, f, protocol=2)


def walk_all_files_with_suffix(dir,
                               suffixs=('.jpg', '.png', '.jpeg'),
                               sort=False):
    paths = []
    names = []
    assert os.path.isdir(dir), '%s is not a valid directory' % dir
    for root, _, fnames in sorted(os.walk(dir)):
        for fname in fnames:
            path = os.path.join(root, fname)
            if os.path.splitext(fname)[1].lower() in suffixs:
                paths.append(path)
                if not sort:
                    names.append(fname)

    if sort:
        paths.sort()
        for i in range(len(paths)):
            names.append(os.path.basename(paths[i]))
    return len(names), names, paths


def get_dirs(root_dir):
    dir_paths = []
    dir_names = []
    for lists in os.listdir(root_dir):
        path = os.path.join(root_dir, lists)
        if os.path.isdir(path):
            dir_paths.append(path)
            dir_names.append(os.path.basename(path))
    dir_names.sort()
    dir_paths.sort()
    return len(dir_names), dir_names, dir_paths


def get_leave_dirs(root_dir):
    leave_dirs = []
    for dirpath, dirnames, filenames in os.walk(root_dir):
        if not dirnames: 
            leave_dirs.append(dirpath)
    return leave_dirs


def merge_pkl_dict(paths):
    dicts = []
    for i in range(len(paths)):
        dicts.append(pickle_load(paths[i]))
    for i in range(len(paths) - 1):
        dicts[0].update(dicts[i + 1])
    return dicts[0]