Example of image resize step with support for archived inputs and outputs without using valohai-utils

without-valohai-utils.py

import glob
import os
import argparse
import json
import shutil
import tempfile
import zipfile
from PIL import Image

parser = argparse.ArgumentParser()
parser.add_argument('--width', default=640, type=int, help='Image width')
parser.add_argument('--height', default=480, type=int, help='Image height')
parser.add_argument('--images_path', default="./inputs.zip", type=str, help='Images path')
parser.add_argument('--output_path', default="./resized.zip", type=str, help='Images path')
args = parser.parse_args()


def is_running_in_valohai():
    return bool(os.environ.get("VH_JOB_ID"))


def resize_image(in_path, out_path, width, height):
    image = Image.open(in_path)
    print(json.dumps({
        "from_width": image.size[0],
        "from_height": image.size[0],
        "to_width": width,
        "to_height": height,
    }))
    new_image = image.resize((width, height))
    new_image.save(out_path)


if __name__ == '__main__':
    images_path = "/valohai/inputs/images/*.zip" if is_running_in_valohai() else args.images_path
    output_path = "/valohai/outputs/resized" if is_running_in_valohai() else args.output_path

    if is_running_in_valohai():
        input_path = "/valohai/inputs/images/*.zip"
    else:
        input_path = "/tmp/*.zip"

    indir = tempfile.mkdtemp()
    outdir = tempfile.mkdtemp()

    for archive_path in glob.glob(input_path):
        with zipfile.ZipFile(archive_path, 'r') as zip_ref:
            zip_ref.extractall(indir)
            for image_path in glob.glob(os.path.join(indir, "*")):
                resize_image(
                    in_path=image_path,
                    out_path=os.path.join(outdir, image_path),
                    width=args.width,
                    height=args.height,
                )

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    shutil.make_archive(os.path.splitext(output_path)[0], 'zip', outdir)
    shutil.rmtree(indir)
    shutil.rmtree(outdir)