# Copyright 2022 The MediaPipe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Test utilities for model maker vision module."""

from typing import Collection

import numpy as np
import tensorflow as tf


def fill_image(rgb: Collection[int], image_size: int):
  """Test helper function to create images.

  Args:
    rgb: A tuple or array of rgb values in [r, g, b] format
    image_size: Int specifying the edge of the square image

  Returns:
    Numpy array of shape (image_size, image_size, 3) filled with the rgb color
  """
  r, g, b = rgb
  return np.broadcast_to(
      np.array([[[r, g, b]]], dtype=np.uint8), shape=(image_size, image_size, 3)
  )


def write_filled_jpeg_file(path: str, rgb: Collection[int], image_size: int):
  """Writes an image to a file path.

  Args:
    path: location to write the image
    rgb: A tuple or array of rgb values in [r, g, b] format
    image_size: Int specifying the edge of the square image
  """
  tf.keras.preprocessing.image.save_img(
      path=path,
      x=fill_image(rgb, image_size),
      data_format='channels_last',
      file_format='jpeg',
  )
