Remove or move NIfTI images in a directory that are not the most common size

prune_on_dim.py

#!/usr/bin/env python
# prune files from a directory that are different dimensions than the mode

import argparse
from collections import Counter
from glob import glob
import os
import shutil
import sys

import ants


def arg_parser():
    parser = argparse.ArgumentParser(description="remove files with dimension that is different from the mode in a directory")
    parser.add_argument('img_dir', type=str)
    parser.add_argument('-m', '--move_dir', type=str, default=None)
    return parser


def split_filename(filepath):
    path = os.path.dirname(filepath)
    filename = os.path.basename(filepath)
    base, ext = os.path.splitext(filename)
    if ext == '.gz':
        base, ext2 = os.path.splitext(base)
        ext = ext2 + ext
    return path, base, ext


def main():
    try:
        args = arg_parser().parse_args()
        shapes = {}
        for fn in glob(os.path.join(args.img_dir, '*.nii.gz')):
            img = ants.image_read(fn)
            print(f'{fn}. Spacing: {img.spacing}, Shape: {img.shape}')
            shapes[fn] = img.shape
        c = Counter(shapes.values())
        mode = c.most_common()[0][0]
        for fn in glob(os.path.join(args.img_dir, '*.nii.gz')):
            if shapes[fn] != mode:
                if args.move_dir is None:
                    print(f'Removing file: {fn} since it has shape {shapes[fn]} != {mode}')
                    os.remove(fn)
                else:
                    print(f'Moving file to {args.move_dir}: {fn} since it has shape {shapes[fn]} != {mode}')
                    _, base, ext = split_filename(fn)
                    shutil.move(fn, os.path.join(args.move_dir, base + ext))
        return 0
    except Exception as e:
        print(e)
        return 1


if __name__ == "__main__":
    sys.exit(main())