#!/usr/bin/env python3

# Copyright 2022 The IREE Authors
#
# Licensed under the Apache License v2.0 with LLVM Exceptions.
# See https://llvm.org/LICENSE.txt for license information.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

# This script converts an image into a 28x28 grayscale f32 buffer.
# Usage:
#   cat image.png | python3 convert_image.py > converted_image.bin

import sys
from PIL import Image
import numpy as np

# Read image from stdin (in any format supported by PIL).
with Image.open(sys.stdin.buffer) as color_img:
  # Resize to 28x28, matching what the program expects.
  resized_color_img = color_img.resize((28, 28))
  # Convert to grayscale.
  grayscale_img = resized_color_img.convert('L')
  # Rescale to a float32 in range [0.0, 1.0].
  grayscale_arr = np.array(grayscale_img)
  grayscale_arr_f32 = grayscale_arr.astype(np.float32) / 255.0
  # Write bytes back out to stdout.
  sys.stdout.buffer.write(grayscale_arr_f32.tobytes())