cropcsv.py

import argparse
import csv
import os
import subprocess
from datetime import datetime
from multiprocessing import Pool, cpu_count
from pathlib import Path
from shutil import move

from lib import FullPaths, is_dir, __VERSION__, __AUTHOR__, __YEAR__, str2bool

__DESCRIPTION__ = "Create cropped images from CSV data generated by gen.py with `convert` (ImageMagick)"
__EPILOG__ = "%(prog)s v{0} (c) {1} {2}-".format(__VERSION__, __AUTHOR__, __YEAR__)
__EXAMPLES__ = [
]


def convert_multi(args):
    return convert(*args)


def convert(path: Path, width, height, startX, startY, debug=False) -> str:
    newname = Path(os.path.join(path.parent, "crop." + path.name)).absolute()

    cmd = [
        "convert",
        str(path),
    ]

    if debug:
        cmd.extend([
            "-fill", "none",
            "-stroke", "red",
            "-draw", f"rectangle {startX},{startY} {startX + width},{height}"
        ])
    else:
        cmd.extend(["-crop", f"{width}x{height}+{startX}+{startY}"])

    # output filename
    cmd.append(os.path.join(newname.parent, '.tmp_' + newname.name))

    print(" ".join(cmd))

    result = subprocess.run(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        universal_newlines=True,
    )

    if result.returncode != 0:
        raise RuntimeError(result.stderr)

    if result.stderr != "":
        raise RuntimeError(result.stderr)

    movedname = move(
        os.path.join(newname.parent, '.tmp_' + newname.name),
        os.path.join(newname.parent, newname.name)
    )

    return movedname


def fmt(src: list) -> (int, int, int, int, int, int, int, int, str):
    for i in range(8):
        src[i] = int(src[i])

    return tuple(src)


if __name__ == '__main__':
    parser: argparse.ArgumentParser = argparse.ArgumentParser(
        description=__DESCRIPTION__,
        epilog=__EPILOG__,
        usage=os.linesep.join(__EXAMPLES__),
    )

    parser.add_argument(
        '--verbose', '-v',
        action='count',
        required=False,
        default=0,
        dest='verbose',
        help="Be verbose. -vvv..v Be more verbose.",
    )

    parser.add_argument(
        '--csv',
        default="crop.csv",
        type=str,
        dest='fname',
        required=False,
        help='CSV filename containing cropping data',
    )

    parser.add_argument(
        '--debug',
        default=False,
        type=str2bool,
        const=True,
        dest='debug',
        required=False,
        nargs='?',
        help='Enable debug mode (draws rectangles over cropping area instead of cropping)',
    )

    parser.add_argument(
        '--stop',
        default=0,
        type=int,
        dest='stop',
        required=False,
        help='Stop at this frame, 0 = no limit',
    )

    parser.add_argument(
        action=FullPaths,
        type=is_dir,
        dest='dir',
        help='Directory containing CSV and images',
    )

    args = parser.parse_args()

    lines: list = []

    print("Reading CSV...")
    with open(os.path.join(args.dir, args.fname), 'r', encoding='utf-8') as f:
        rdr = csv.reader(f)
        for idx, i in enumerate(rdr):
            if idx == 0:
                # First line has headers, skip
                continue

            if args.stop != 0 and idx == args.stop:
                # Stop at given frame
                break

            startX, startY, endX, endY, width, height, owidth, oheight, fname = fmt(i)

            if (startX + width) > owidth:
                startX = owidth - width
            if (startY + height) > oheight:
                startY = oheight - height

            lines.append((
                Path(os.path.join(args.dir, fname)).absolute(),
                width,
                height,
                startX,
                startY,
                args.debug
            ))

    cpus: int = cpu_count() - 1
    if cpus < 1:
        cpus = 1

    print(f"Processing with {cpus} CPUs...")
    now = datetime.now()
    with Pool(processes=cpus) as pool:
        results = pool.map(convert_multi, lines)
    print("Took", str(datetime.now() - now))
    print("Done.")