|
|
|
|
|
import argparse |
|
import numpy as np |
|
import sys |
|
import cv2 |
|
|
|
import tritonclient.grpc as grpcclient |
|
from tritonclient.utils import InferenceServerException |
|
|
|
from processing import preprocess, postprocess |
|
from render import render_box, render_filled_box, get_text_size, render_text, RAND_COLORS |
|
from labels import COCOLabels |
|
|
|
INPUT_NAMES = ["images"] |
|
OUTPUT_NAMES = ["num_dets", "det_boxes", "det_scores", "det_classes"] |
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser() |
|
parser.add_argument('mode', |
|
choices=['dummy', 'image', 'video'], |
|
default='dummy', |
|
help='Run mode. \'dummy\' will send an emtpy buffer to the server to test if inference works. \'image\' will process an image. \'video\' will process a video.') |
|
parser.add_argument('input', |
|
type=str, |
|
nargs='?', |
|
help='Input file to load from in image or video mode') |
|
parser.add_argument('-m', |
|
'--model', |
|
type=str, |
|
required=False, |
|
default='yolov7', |
|
help='Inference model name, default yolov7') |
|
parser.add_argument('--width', |
|
type=int, |
|
required=False, |
|
default=640, |
|
help='Inference model input width, default 640') |
|
parser.add_argument('--height', |
|
type=int, |
|
required=False, |
|
default=640, |
|
help='Inference model input height, default 640') |
|
parser.add_argument('-u', |
|
'--url', |
|
type=str, |
|
required=False, |
|
default='localhost:8001', |
|
help='Inference server URL, default localhost:8001') |
|
parser.add_argument('-o', |
|
'--out', |
|
type=str, |
|
required=False, |
|
default='', |
|
help='Write output into file instead of displaying it') |
|
parser.add_argument('-f', |
|
'--fps', |
|
type=float, |
|
required=False, |
|
default=24.0, |
|
help='Video output fps, default 24.0 FPS') |
|
parser.add_argument('-i', |
|
'--model-info', |
|
action="store_true", |
|
required=False, |
|
default=False, |
|
help='Print model status, configuration and statistics') |
|
parser.add_argument('-v', |
|
'--verbose', |
|
action="store_true", |
|
required=False, |
|
default=False, |
|
help='Enable verbose client output') |
|
parser.add_argument('-t', |
|
'--client-timeout', |
|
type=float, |
|
required=False, |
|
default=None, |
|
help='Client timeout in seconds, default no timeout') |
|
parser.add_argument('-s', |
|
'--ssl', |
|
action="store_true", |
|
required=False, |
|
default=False, |
|
help='Enable SSL encrypted channel to the server') |
|
parser.add_argument('-r', |
|
'--root-certificates', |
|
type=str, |
|
required=False, |
|
default=None, |
|
help='File holding PEM-encoded root certificates, default none') |
|
parser.add_argument('-p', |
|
'--private-key', |
|
type=str, |
|
required=False, |
|
default=None, |
|
help='File holding PEM-encoded private key, default is none') |
|
parser.add_argument('-x', |
|
'--certificate-chain', |
|
type=str, |
|
required=False, |
|
default=None, |
|
help='File holding PEM-encoded certicate chain default is none') |
|
|
|
FLAGS = parser.parse_args() |
|
|
|
|
|
try: |
|
triton_client = grpcclient.InferenceServerClient( |
|
url=FLAGS.url, |
|
verbose=FLAGS.verbose, |
|
ssl=FLAGS.ssl, |
|
root_certificates=FLAGS.root_certificates, |
|
private_key=FLAGS.private_key, |
|
certificate_chain=FLAGS.certificate_chain) |
|
except Exception as e: |
|
print("context creation failed: " + str(e)) |
|
sys.exit() |
|
|
|
|
|
if not triton_client.is_server_live(): |
|
print("FAILED : is_server_live") |
|
sys.exit(1) |
|
|
|
if not triton_client.is_server_ready(): |
|
print("FAILED : is_server_ready") |
|
sys.exit(1) |
|
|
|
if not triton_client.is_model_ready(FLAGS.model): |
|
print("FAILED : is_model_ready") |
|
sys.exit(1) |
|
|
|
if FLAGS.model_info: |
|
|
|
try: |
|
metadata = triton_client.get_model_metadata(FLAGS.model) |
|
print(metadata) |
|
except InferenceServerException as ex: |
|
if "Request for unknown model" not in ex.message(): |
|
print("FAILED : get_model_metadata") |
|
print("Got: {}".format(ex.message())) |
|
sys.exit(1) |
|
else: |
|
print("FAILED : get_model_metadata") |
|
sys.exit(1) |
|
|
|
|
|
try: |
|
config = triton_client.get_model_config(FLAGS.model) |
|
if not (config.config.name == FLAGS.model): |
|
print("FAILED: get_model_config") |
|
sys.exit(1) |
|
print(config) |
|
except InferenceServerException as ex: |
|
print("FAILED : get_model_config") |
|
print("Got: {}".format(ex.message())) |
|
sys.exit(1) |
|
|
|
|
|
if FLAGS.mode == 'dummy': |
|
print("Running in 'dummy' mode") |
|
print("Creating emtpy buffer filled with ones...") |
|
inputs = [] |
|
outputs = [] |
|
inputs.append(grpcclient.InferInput(INPUT_NAMES[0], [1, 3, FLAGS.width, FLAGS.height], "FP32")) |
|
inputs[0].set_data_from_numpy(np.ones(shape=(1, 3, FLAGS.width, FLAGS.height), dtype=np.float32)) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[0])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[1])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[2])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[3])) |
|
|
|
print("Invoking inference...") |
|
results = triton_client.infer(model_name=FLAGS.model, |
|
inputs=inputs, |
|
outputs=outputs, |
|
client_timeout=FLAGS.client_timeout) |
|
if FLAGS.model_info: |
|
statistics = triton_client.get_inference_statistics(model_name=FLAGS.model) |
|
if len(statistics.model_stats) != 1: |
|
print("FAILED: get_inference_statistics") |
|
sys.exit(1) |
|
print(statistics) |
|
print("Done") |
|
|
|
for output in OUTPUT_NAMES: |
|
result = results.as_numpy(output) |
|
print(f"Received result buffer \"{output}\" of size {result.shape}") |
|
print(f"Naive buffer sum: {np.sum(result)}") |
|
|
|
|
|
if FLAGS.mode == 'image': |
|
print("Running in 'image' mode") |
|
if not FLAGS.input: |
|
print("FAILED: no input image") |
|
sys.exit(1) |
|
|
|
inputs = [] |
|
outputs = [] |
|
inputs.append(grpcclient.InferInput(INPUT_NAMES[0], [1, 3, FLAGS.width, FLAGS.height], "FP32")) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[0])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[1])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[2])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[3])) |
|
|
|
print("Creating buffer from image file...") |
|
input_image = cv2.imread(str(FLAGS.input)) |
|
if input_image is None: |
|
print(f"FAILED: could not load input image {str(FLAGS.input)}") |
|
sys.exit(1) |
|
input_image_buffer = preprocess(input_image, [FLAGS.width, FLAGS.height]) |
|
input_image_buffer = np.expand_dims(input_image_buffer, axis=0) |
|
|
|
inputs[0].set_data_from_numpy(input_image_buffer) |
|
|
|
print("Invoking inference...") |
|
results = triton_client.infer(model_name=FLAGS.model, |
|
inputs=inputs, |
|
outputs=outputs, |
|
client_timeout=FLAGS.client_timeout) |
|
if FLAGS.model_info: |
|
statistics = triton_client.get_inference_statistics(model_name=FLAGS.model) |
|
if len(statistics.model_stats) != 1: |
|
print("FAILED: get_inference_statistics") |
|
sys.exit(1) |
|
print(statistics) |
|
print("Done") |
|
|
|
for output in OUTPUT_NAMES: |
|
result = results.as_numpy(output) |
|
print(f"Received result buffer \"{output}\" of size {result.shape}") |
|
print(f"Naive buffer sum: {np.sum(result)}") |
|
|
|
num_dets = results.as_numpy(OUTPUT_NAMES[0]) |
|
det_boxes = results.as_numpy(OUTPUT_NAMES[1]) |
|
det_scores = results.as_numpy(OUTPUT_NAMES[2]) |
|
det_classes = results.as_numpy(OUTPUT_NAMES[3]) |
|
detected_objects = postprocess(num_dets, det_boxes, det_scores, det_classes, input_image.shape[1], input_image.shape[0], [FLAGS.width, FLAGS.height]) |
|
print(f"Detected objects: {len(detected_objects)}") |
|
|
|
for box in detected_objects: |
|
print(f"{COCOLabels(box.classID).name}: {box.confidence}") |
|
input_image = render_box(input_image, box.box(), color=tuple(RAND_COLORS[box.classID % 64].tolist())) |
|
size = get_text_size(input_image, f"{COCOLabels(box.classID).name}: {box.confidence:.2f}", normalised_scaling=0.6) |
|
input_image = render_filled_box(input_image, (box.x1 - 3, box.y1 - 3, box.x1 + size[0], box.y1 + size[1]), color=(220, 220, 220)) |
|
input_image = render_text(input_image, f"{COCOLabels(box.classID).name}: {box.confidence:.2f}", (box.x1, box.y1), color=(30, 30, 30), normalised_scaling=0.5) |
|
|
|
if FLAGS.out: |
|
cv2.imwrite(FLAGS.out, input_image) |
|
print(f"Saved result to {FLAGS.out}") |
|
else: |
|
cv2.imshow('image', input_image) |
|
cv2.waitKey(0) |
|
cv2.destroyAllWindows() |
|
|
|
|
|
if FLAGS.mode == 'video': |
|
print("Running in 'video' mode") |
|
if not FLAGS.input: |
|
print("FAILED: no input video") |
|
sys.exit(1) |
|
|
|
inputs = [] |
|
outputs = [] |
|
inputs.append(grpcclient.InferInput(INPUT_NAMES[0], [1, 3, FLAGS.width, FLAGS.height], "FP32")) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[0])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[1])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[2])) |
|
outputs.append(grpcclient.InferRequestedOutput(OUTPUT_NAMES[3])) |
|
|
|
print("Opening input video stream...") |
|
cap = cv2.VideoCapture(FLAGS.input) |
|
if not cap.isOpened(): |
|
print(f"FAILED: cannot open video {FLAGS.input}") |
|
sys.exit(1) |
|
|
|
counter = 0 |
|
out = None |
|
print("Invoking inference...") |
|
while True: |
|
ret, frame = cap.read() |
|
if not ret: |
|
print("failed to fetch next frame") |
|
break |
|
|
|
if counter == 0 and FLAGS.out: |
|
print("Opening output video stream...") |
|
fourcc = cv2.VideoWriter_fourcc('M', 'P', '4', 'V') |
|
out = cv2.VideoWriter(FLAGS.out, fourcc, FLAGS.fps, (frame.shape[1], frame.shape[0])) |
|
|
|
input_image_buffer = preprocess(frame, [FLAGS.width, FLAGS.height]) |
|
input_image_buffer = np.expand_dims(input_image_buffer, axis=0) |
|
|
|
inputs[0].set_data_from_numpy(input_image_buffer) |
|
|
|
results = triton_client.infer(model_name=FLAGS.model, |
|
inputs=inputs, |
|
outputs=outputs, |
|
client_timeout=FLAGS.client_timeout) |
|
|
|
num_dets = results.as_numpy("num_dets") |
|
det_boxes = results.as_numpy("det_boxes") |
|
det_scores = results.as_numpy("det_scores") |
|
det_classes = results.as_numpy("det_classes") |
|
detected_objects = postprocess(num_dets, det_boxes, det_scores, det_classes, frame.shape[1], frame.shape[0], [FLAGS.width, FLAGS.height]) |
|
print(f"Frame {counter}: {len(detected_objects)} objects") |
|
counter += 1 |
|
|
|
for box in detected_objects: |
|
print(f"{COCOLabels(box.classID).name}: {box.confidence}") |
|
frame = render_box(frame, box.box(), color=tuple(RAND_COLORS[box.classID % 64].tolist())) |
|
size = get_text_size(frame, f"{COCOLabels(box.classID).name}: {box.confidence:.2f}", normalised_scaling=0.6) |
|
frame = render_filled_box(frame, (box.x1 - 3, box.y1 - 3, box.x1 + size[0], box.y1 + size[1]), color=(220, 220, 220)) |
|
frame = render_text(frame, f"{COCOLabels(box.classID).name}: {box.confidence:.2f}", (box.x1, box.y1), color=(30, 30, 30), normalised_scaling=0.5) |
|
|
|
if FLAGS.out: |
|
out.write(frame) |
|
else: |
|
cv2.imshow('image', frame) |
|
if cv2.waitKey(1) == ord('q'): |
|
break |
|
|
|
if FLAGS.model_info: |
|
statistics = triton_client.get_inference_statistics(model_name=FLAGS.model) |
|
if len(statistics.model_stats) != 1: |
|
print("FAILED: get_inference_statistics") |
|
sys.exit(1) |
|
print(statistics) |
|
print("Done") |
|
|
|
cap.release() |
|
if FLAGS.out: |
|
out.release() |
|
else: |
|
cv2.destroyAllWindows() |
|
|
