inference_api.py

import requests
import argparse
import sys
import base64

# Default examples
# api = "http://localhost:8080/2015-03-31/functions/function/invocations"
# file = "./tests/data/boats.jpg"


def arg_parser():
    """Parse arguments"""

    # Create an ArgumentParser object
    parser = argparse.ArgumentParser(description='Object detection inference via API call')
    # Add arguments
    parser.add_argument('--api', type=str, help='URL to server API (with endpoint)', required=True)
    parser.add_argument('--file', type=str, help='Path to the input image file', required=True)
    parser.add_argument('--model', type=str, \
                        choices=['detr-resnet-50', 'detr-resnet-101', 'yolos-tiny', 'yolos-small'], \
                        help='Model type', \
                        required=False)
    parser.add_argument('-v', '--verbose', action='store_true', help='Increase output verbosity')
    return parser


def main(args=None):
    """Main function"""

    args = arg_parser().parse_args(args)
    # Use the arguments
    if args.verbose:
        print(f'Input file: {args.file}')

    # Retrieve model type
    if args.model:
        model_name = args.model
    else:
        model_name = ""

    # Load image
    with open(args.file, 'rb') as image_file:
        image_data = image_file.read()

    # Encode the image data in base64
    encoded_image = base64.b64encode(image_data).decode('utf-8')

    # Prepare the payload
    payload = {
        'body': encoded_image,
        'isBase64Encoded': True,
        'model': model_name,
    }

    # Send request to API
    # Option 'files': A dictionary of files to send to the specified url
    # response = requests.post(args.api, files={'image': image_data})
    # Option 'json': A JSON object to send to the specified url
    response = requests.post(args.api, json = payload)

    if response.status_code == 200:
        print('Detection Results:')
        # Process the response
        # processed_data = json.loads(response.content)
        # print('processed_data', processed_data)
        results = response.json()
        print("results: ", results)
    else:
        print(f"Error: {response.status_code}")
        print(response.json())


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))