Spaces:
Sleeping
Sleeping
Clement Vachet
commited on
Commit
·
b58c566
1
Parent(s):
71bec2b
Add model type as new argument
Browse files- 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
|