paddlers/tools/prepare_dataset/common.py

import argparse
import os
import os.path as osp
from glob import glob
from itertools import count
from functools import partial
from concurrent.futures import ThreadPoolExecutor

from skimage.io import imread, imsave
from tqdm import tqdm


def get_default_parser():
    """
    Get argument parser with commonly used options.
    
    Returns:
        argparse.ArgumentParser: Argument parser with the following arguments:
            --in_dataset_dir: Input dataset directory.
            --out_dataset_dir: Output dataset directory.
    """

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--in_dataset_dir',
        type=str,
        required=True,
        help="Input dataset directory.")
    parser.add_argument(
        '--out_dataset_dir', type=str, help="Output dataset directory.")
    return parser


def add_crop_options(parser):
    """
    Add patch cropping related arguments to an argument parser. The parser will be
        modified in place.
    
    Args:
        parser (argparse.ArgumentParser): Argument parser.
    
    Returns:
        argparse.ArgumentParser: Argument parser with the following arguments:
            --crop_size: Size of cropped patches.
            --crop_stride: Stride of sliding windows when cropping patches.
    """

    parser.add_argument(
        '--crop_size', type=int, help="Size of cropped patches.")
    parser.add_argument(
        '--crop_stride',
        type=int,
        help="Stride of sliding windows when cropping patches. `crop_size` will be used only if `crop_size` is not None.",
    )
    return parser


def crop_and_save(path, out_subdir, crop_size, stride):
    name, ext = osp.splitext(osp.basename(path))
    out_subsubdir = osp.join(out_subdir, name)
    if not osp.exists(out_subsubdir):
        os.makedirs(out_subsubdir)
    img = imread(path)
    w, h = img.shape[:2]
    counter = count()
    for i in range(0, h - crop_size + 1, stride):
        for j in range(0, w - crop_size + 1, stride):
            imsave(
                osp.join(out_subsubdir, '{}_{}{}'.format(name,
                                                         next(counter), ext)),
                img[i:i + crop_size, j:j + crop_size],
                check_contrast=False)


def crop_patches(crop_size,
                 stride,
                 data_dir,
                 out_dir,
                 subsets=('train', 'val', 'test'),
                 subdirs=('A', 'B', 'label'),
                 glob_pattern='*',
                 max_workers=0):
    """
    Crop patches from images in specific directories.
    
    Args:
        crop_size (int): Height and width of the cropped patches will be `crop_size`.
        stride (int): Stride of sliding windows when cropping patches.
        data_dir (str): Root directory of the dataset that contains the input images.
        out_dir (str): Directory to save the cropped patches.
        subsets (tuple|list|None, optional): List or tuple of names of subdirectories 
            or None. Images to be cropped should be stored in `data_dir/subset/subdir/` 
            or `data_dir/subdir/` (when `subsets` is set to None), where `subset` is an 
            element of `subsets`. Defaults to ('train', 'val', 'test').
        subdirs (tuple|list, optional): List or tuple of names of subdirectories. Images 
            to be cropped should be stored in `data_dir/subset/subdir/` or 
            `data_dir/subdir/` (when `subsets` is set to None), where `subdir` is an 
            element of `subdirs`. Defaults to ('A', 'B', 'label').
        glob_pattern (str, optional): Glob pattern used to match image files. 
            Defaults to '*', which matches arbitrary file. 
        max_workers (int, optional): Number of worker threads to perform the cropping 
            operation. Deafults to 0.
    """

    if max_workers < 0:
        raise ValueError("`max_workers` must be a non-negative integer!")

    if subset is None:
        subsets = ('', )

    if max_workers == 0:
        for subset in subsets:
            for subdir in subdirs:
                paths = glob(
                    osp.join(data_dir, subset, subdir, glob_pattern),
                    recursive=True)
                out_subdir = osp.join(out_dir, subset, subdir)
                for p in tqdm(paths):
                    crop_and_save(
                        p,
                        out_subdir=out_subdir,
                        crop_size=crop_size,
                        stride=stride)
    else:
        # Concurrently crop image patches
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            for subset in subsets:
                for subdir in subdirs:
                    paths = glob(
                        osp.join(data_dir, subset, subdir, glob_pattern),
                        recursive=True)
                    out_subdir = osp.join(out_dir, subset, subdir)
                    for _ in tqdm(
                            executor.map(partial(
                                crop_and_save,
                                out_subdir=out_subdir,
                                crop_size=crop_size,
                                stride=stride),
                                         paths),
                            total=len(paths)):
                        pass


def get_path_tuples(*dirs, glob_pattern='*', data_dir=None):
    """
    Get tuples of image paths. Each tuple corresponds to a sample in the dataset.
    
    Args:
        *dirs (str): Directories that contains the images.
        glob_pattern (str, optional): Glob pattern used to match image files. 
            Defaults to '*', which matches arbitrary file. 
        data_dir (str|None, optional): Root directory of the dataset that 
            contains the images. If not None, `data_dir` will be used to 
            determine relative paths of images. Defaults to None.
    
    Returns:
        list[tuple]: For directories with the following structure:
            ├── img  
            │   ├── im1.png
            │   ├── im2.png
            │   └── im3.png
            │
            ├── mask
            │   ├── im1.png
            │   ├── im2.png
            │   └── im3.png
            └── ...

        `get_path_tuples('img', 'mask', '*.png')` will return list of tuples:
            [('img/im1.png', 'mask/im1.png'), ('img/im2.png', 'mask/im2.png'), ('img/im3.png', 'mask/im3.png')]
    """

    all_paths = []
    for dir_ in dirs:
        paths = glob(osp.join(dir_, glob_pattern), recursive=True)
        paths = sorted(paths)
        if data_dir is not None:
            paths = [osp.relpath(p, data_dir) for p in paths]
        all_paths.append(paths)
    all_paths = list(zip(*all_paths))
    return all_paths


def create_file_list(file_list, path_tuples, sep=' '):
    """
    Create file list.
    
    Args:
        file_list (str): Path of file list to create.
        path_tuples (list[tuple]): See get_path_tuples().
        sep (str, optional): Delimiter to use when writing lines to file list. 
            Defaults to ' '.
    """

    with open(file_list, 'w') as f:
        for tup in path_tuples:
            line = sep.join(tup)
            f.write(line + '\n')


def link_dataset(src, dst):
    """
    Make a symbolic link to a dataset.
    
    Args:
        src (str): Path of the original dataset.
        dst (str): Path of the symbolic link.
    """

    if osp.exists(dst) and not osp.isdir(dst):
        raise ValueError(f"{dst} exists and is not a directory.")
    elif not osp.exists(dst):
        os.makedirs(dst)
    name = osp.basename(osp.normpath(src))
    os.symlink(src, osp.join(dst, name), target_is_directory=True)