test_post_request.py

import os
import cv2
import base64
import logging
import argparse
import requests
import numpy as np

from skimage.io import imread, imsave


def send_image_post_request(ARGS: argparse.Namespace) -> None:
    logging.basicConfig(level=logging.INFO)
    list_test_images = os.listdir(ARGS.dir_test_images)

    for file_image in list_test_images:
        target_file_image = file_image
        file_image = os.path.join(ARGS.dir_test_images, file_image)
        files = {"image_file": (file_image, open(file_image, "rb"))}

        # if the deployment is on local machine
        # response = requests.post(
        #     "http://127.0.0.1:7860/generate",
        #     files=files,
        # )

        # if the deployment is on hugging face
        response = requests.post(
            "https://abhishekrs4-image-colourization-cgan.hf.space/generate",
            files=files,
        )

        # logging.info(response.json())
        encoded_image = response.json()["encoded_image"]
        image_str = base64.b64decode(encoded_image)
        image_dec = cv2.imdecode(np.frombuffer(image_str, np.uint8), cv2.IMREAD_COLOR)
        logging.info(f"image_sum: {np.sum(image_dec)}")
        imsave(os.path.join(ARGS.dir_gen_images, target_file_image), image_dec)
    return


def main() -> None:
    dir_test_images = "./sample_test_images/"
    dir_gen_images = "./generated_color_images/"
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument(
        "--dir_test_images",
        default=dir_test_images,
        type=str,
        help="full directory path to dataset containing test images",
    )
    parser.add_argument(
        "--dir_gen_images",
        default=dir_gen_images,
        type=str,
        help="full directory path to directory where generated color images needs to be saved",
    )
    ARGS, unparsed = parser.parse_known_args()
    send_image_post_request(ARGS)
    return


if __name__ == "__main__":
    main()