convert_images.py

import argparse
import glob
import json
import os

import torch

from utils import (
    preprocess_image_and_quantize,
    torch_type_to_sgdk_type,
    write_array,
)


def generate(input: str, output: str, scale: float = 1.0, zero_point: int = 0) -> None:
    images = []
    if os.path.isdir(input):
        for file in glob.glob(os.path.join(input, "*.jpg")):
            image = preprocess_image_and_quantize(file, scale, zero_point).squeeze(0)
            images.append(image)
    if os.path.isfile(input):
        image = preprocess_image_and_quantize(input, scale, zero_point).squeeze(0)
        images.append(image)

    size = len(images)
    images = torch.cat(images, dim=0)
    with open(output, "w") as f:
        f.write(f"/* Auto-generated from {input} */\n")
        f.write("#ifndef _MNIST_H_\n")
        f.write("#define _MNIST_H_\n\n")
        f.write("#include <genesis.h>\n\n")
        f.write("#define CLASSES_COUNT 10\n")

        f.write("#define IMG_H 28\n")
        f.write("#define IMG_W 28\n\n")

        f.write(f"#define IMG_COUNT {size}\n\n")

        write_array(
            f,
            "images",
            images.tolist(),
            [size, 28, 28],
            torch_type_to_sgdk_type(str(images.dtype)),
        )

        f.write("#endif")


def _get_model_input_qparams(model_path: str) -> tuple[float, int]:
    with open(model_path, "r") as f:
        model = json.load(f)
    return model["fc1_act_scale"]["data"][0], model["fc1_act_zero_point"]["data"][0]


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="Image2Header",
        description="Converts 28x28 MNIST images to C header file",
    )
    parser.add_argument(
        "-i",
        "--input",
        type=str,
        default="./data/MNIST/images",
        help="Path to JSON weights file",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default="./inc/mnist.h",
        help="Output C header file path",
    )
    parser.add_argument(
        "--model",
        type=str,
        default="./data/models/quantized_weights.json",
        help="Path to JSON quantized weights file",
    )

    args = parser.parse_args()

    scale, zero_point = _get_model_input_qparams(args.model)
    generate(input=args.input, output=args.output, scale=scale, zero_point=zero_point)