object_detection_lambda / inference_api.py
Clement Vachet
Add model type as new argument
b58c566
raw
history blame
2.29 kB
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:]))