stitch-videoshots.py

#! /usr/bin/env python3
"""Combine multiple screenshots with the same shape.

Usage
-----
    1. Specify a list of screenshots.
        stitch-videoshots screenshot1 screenshot2 screenshot3
    2. Specify a dictionary contains screenshots.
        stitch-videoshots -d ./mypath [--sort]

"""

import argparse
from pathlib import Path
from PIL import Image
from datetime import datetime


def stitch_imgs(imgs, savename):
    """Stitch all input images together.

    Parameters
    ----------
    imgs : list of sfr or Path
        Path to the input screenshots.
    savename : str
        Pre-defined name to save the stitched image.
    """
    imgs = [Image.open(i) for i in imgs]

    width, height = imgs[0].size
    stitched_img = Image.new('RGB', (width, height * len(imgs)))
    height_offset = 0

    for img in imgs:
        stitched_img.paste(img, (0, height_offset))
        height_offset += height

    stitched_img.save(savename)
    print(f'Saved stitched image to ./{savename} successfully ...')


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('screenshots', nargs='*',
                        type=str,
                        help='A list of screenshots you\'d like to stitch together')
    parser.add_argument('-d', '--dirpath',
                        type=str,
                        help='The directory contains your screenshots')
    parser.add_argument('--sort',
                        action='store_true',
                        help='Set to true if you want stitch screenshots based on the filename alphabetical order')
    args = parser.parse_args()

    if args.screenshots:
        screenshots = [Path(i) for i in (sorted(args.screenshots)
                                         if args.sort else args.screenshots)]
        filename = screenshots[0].stem
    elif args.dirpath is None:
        raise IOError('Path to a directory or a list of screenshots is expected. '
                      'Check out -h for helping message.')
    else:
        screenshots = [i for i in (sorted(Path(args.dirpath).glob('*'))
                                   if args.sort else Path(args.dirpath).glob('*'))
                       if i.suffix in ['.jpg', '.png', '.jpeg']]
        filename = Path(args.dirpath).resolve().name

    suffix = screenshots[0].suffix

    EXEC_DATE = datetime.strftime(datetime.now(), '%Y-%m-%d-%H-%M-%S')
    savename = filename + '-stitched-' + EXEC_DATE + suffix
    stitch_imgs(screenshots, savename)