Clement Vachet commited on
Commit
b58c566
·
1 Parent(s): 71bec2b

Add model type as new argument

Browse files
Files changed (1) hide show
  1. inference_api.py +13 -1
inference_api.py CHANGED
@@ -16,6 +16,10 @@ def arg_parser():
16
  # Add arguments
17
  parser.add_argument('--api', type=str, help='URL to server API (with endpoint)', required=True)
18
  parser.add_argument('--file', type=str, help='Path to the input image file', required=True)
 
 
 
 
19
  parser.add_argument('-v', '--verbose', action='store_true', help='Increase output verbosity')
20
  return parser
21
 
@@ -28,6 +32,12 @@ def main(args=None):
28
  if args.verbose:
29
  print(f'Input file: {args.file}')
30
 
 
 
 
 
 
 
31
  # Load image
32
  with open(args.file, 'rb') as image_file:
33
  image_data = image_file.read()
@@ -37,7 +47,9 @@ def main(args=None):
37
 
38
  # Prepare the payload
39
  payload = {
40
- 'body': encoded_image
 
 
41
  }
42
 
43
  # Send request to API
 
16
  # Add arguments
17
  parser.add_argument('--api', type=str, help='URL to server API (with endpoint)', required=True)
18
  parser.add_argument('--file', type=str, help='Path to the input image file', required=True)
19
+ parser.add_argument('--model', type=str, \
20
+ choices=['detr-resnet-50', 'detr-resnet-101', 'yolos-tiny', 'yolos-small'], \
21
+ help='Model type', \
22
+ required=False)
23
  parser.add_argument('-v', '--verbose', action='store_true', help='Increase output verbosity')
24
  return parser
25
 
 
32
  if args.verbose:
33
  print(f'Input file: {args.file}')
34
 
35
+ # Retrieve model type
36
+ if args.model:
37
+ model_name = args.model
38
+ else:
39
+ model_name = ""
40
+
41
  # Load image
42
  with open(args.file, 'rb') as image_file:
43
  image_data = image_file.read()
 
47
 
48
  # Prepare the payload
49
  payload = {
50
+ 'body': encoded_image,
51
+ 'isBase64Encoded': True,
52
+ 'model': model_name,
53
  }
54
 
55
  # Send request to API