|
1 | 1 | import os |
2 | | -import cv2 |
3 | 2 | import numpy as np |
4 | | - |
5 | 3 | from tf_proto_utils import InferenceClient |
| 4 | +from vdo_proto_utils import VideoCaptureClient |
6 | 5 |
|
7 | | -# Create the inference client with model and socket from environment variables |
8 | | -inf_client = InferenceClient(os.environ['INFERENCE_HOST']) |
9 | 6 |
|
10 | | -# Create a video stream that fits the model inputs |
11 | | -cap = cv2.VideoCapture(0) |
12 | | -cap.set(cv2.CAP_PROP_FRAME_WIDTH, 300) |
13 | | -cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 300) |
14 | | -cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"RGB3")) |
| 7 | +grpc_socket = os.environ['INFERENCE_HOST'] |
| 8 | +stream_width, stream_height, stream_framerate = (224, 224, 10) |
| 9 | +capture_client = VideoCaptureClient(socket=grpc_socket, |
| 10 | + stream_width=stream_width, |
| 11 | + stream_height=stream_height, |
| 12 | + stream_framerate=stream_framerate) |
| 13 | +inf_client = InferenceClient(grpc_socket) |
15 | 14 |
|
16 | 15 | # The capture -> inference loop |
17 | 16 | while True: |
18 | 17 | # Read an image from the video stream |
19 | | - _, image = cap.read() |
| 18 | + image = capture_client.get_frame() |
20 | 19 |
|
21 | | - # Give the image a batch dimension such that it has the shape (1, H, W, C) |
22 | 20 | image = np.expand_dims(image, axis=0) |
23 | 21 |
|
24 | 22 | # Perform inference with a model specified as an environment variable |
|
0 commit comments