From 98c8868b3921ad07c9f009c94889d2e7d3bca8ce Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Sat, 4 Jul 2020 00:28:56 +1000 Subject: [PATCH 01/11] Implement first prototype for yolov5 example --- examples/onnx/yolov5-videos/README.md | 1 + .../onnx/yolov5-videos/conda-packages.txt | 2 + examples/onnx/yolov5-videos/cortex.yaml | 13 ++ examples/onnx/yolov5-videos/predictor.py | 134 ++++++++++++++++++ examples/onnx/yolov5-videos/requirements.txt | 3 + 5 files changed, 153 insertions(+) create mode 100644 examples/onnx/yolov5-videos/README.md create mode 100644 examples/onnx/yolov5-videos/conda-packages.txt create mode 100644 examples/onnx/yolov5-videos/cortex.yaml create mode 100644 examples/onnx/yolov5-videos/predictor.py create mode 100644 examples/onnx/yolov5-videos/requirements.txt diff --git a/examples/onnx/yolov5-videos/README.md b/examples/onnx/yolov5-videos/README.md new file mode 100644 index 0000000000..a0990367ef --- /dev/null +++ b/examples/onnx/yolov5-videos/README.md @@ -0,0 +1 @@ +TBD diff --git a/examples/onnx/yolov5-videos/conda-packages.txt b/examples/onnx/yolov5-videos/conda-packages.txt new file mode 100644 index 0000000000..49cbb29eaf --- /dev/null +++ b/examples/onnx/yolov5-videos/conda-packages.txt @@ -0,0 +1,2 @@ +conda-forge::ffmpeg +conda-forge::youtube-dl diff --git a/examples/onnx/yolov5-videos/cortex.yaml b/examples/onnx/yolov5-videos/cortex.yaml new file mode 100644 index 0000000000..c6c89a344d --- /dev/null +++ b/examples/onnx/yolov5-videos/cortex.yaml @@ -0,0 +1,13 @@ +# WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version` + +- name: youtube-yolov5 + kind: SyncAPI + predictor: + type: onnx + path: predictor.py + model_path: yolov5s.onnx + config: + iou_threshold: 0.5 + confidence_threshold: 0.3 + compute: + gpu: 1 \ No newline at end of file diff --git a/examples/onnx/yolov5-videos/predictor.py b/examples/onnx/yolov5-videos/predictor.py new file mode 100644 index 0000000000..a2289eeea9 --- /dev/null +++ b/examples/onnx/yolov5-videos/predictor.py @@ -0,0 +1,134 @@ +# WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version` + +import ffmpeg +import youtube_dl +import uuid +from pathlib import Path +import numpy as np +import cv2 +from starlette.responses import FileResponse +from typing import Iterable, Tuple +import os + + +def download_from_youtube(url: str) -> Path: + target = f'{uuid.uuid1()}.mp4' + ydl_opts = {'outtmpl': target, 'format': "worstvideo[vcodec=vp9][height>=480]"} + with youtube_dl.YoutubeDL(ydl_opts) as ydl: + ydl.download([url]) + # we need to glob in case youtube-dl adds suffix + path, = Path().absolute().glob(f'{target}*') + return path + + +def frame_reader(path: Path, size: Tuple[int, int]) -> Iterable[np.ndarray]: + width, height = size + # letterbox frames to fixed size + process = ( + ffmpeg + .input(path) + .filter('scale', size=f'{width}:{height}', force_original_aspect_ratio='decrease') + # Negative values for x and y center the padded video + .filter('pad', height=height, width=width, x=-1, y=-1) + .output('pipe:', format='rawvideo', pix_fmt='rgb24') + .run_async(pipe_stdout=True) + ) + + while True: + in_bytes = process.stdout.read(height * width * 3) + if not in_bytes: + process.wait() + break + frame = np.frombuffer(in_bytes, np.uint8).reshape([height, width, 3]) + yield frame + + +class FrameWriter: + def __init__(self, path, size): + width, height = size + self.process = ( + ffmpeg + .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}') + .output(path, pix_fmt='yuv420p') + .overwrite_output() + .run_async(pipe_stdin=True) + ) + + def write(self, frame): + self.process.stdin.write(frame.astype(np.uint8).tobytes()) + + + def __del__(self): + self.process.stdin.close() + self.process.wait() + + +def nms(dets, scores, thresh): + x1 = dets[:, 0] + y1 = dets[:, 1] + x2 = dets[:, 2] + y2 = dets[:, 3] + + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] # get boxes with more ious first + + keep = [] + while order.size > 0: + i = order[0] # pick maxmum iou box + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) # maximum width + h = np.maximum(0.0, yy2 - yy1 + 1) # maxiumum height + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + inds = np.where(ovr <= thresh)[0] + order = order[inds + 1] + + return keep + + + +class ONNXPredictor: + def __init__(self, onnx_client, config): + self.client = onnx_client + self.config = config + + def postprocess(self, output): + boxes, obj_score, class_scores = np.split(output[0], [4, 5], axis=1) + # convert boxes from x_center, y_center, w, h to xyxy + boxes[:, 0] -= boxes[:, 2] / 2 + boxes[:, 1] -= boxes[:, 3] / 2 + boxes[:, 2] = boxes[:, 2] + boxes[:, 0] + boxes[:, 3] = boxes[:, 3] + boxes[:, 1] + class_id = class_scores.argmax(axis=1) + cls_score = class_scores[np.arange(len(class_scores)), class_id] + confidence = obj_score.squeeze(axis=1) * cls_score + sel = confidence > self.config['confidence_threshold'] + boxes, class_id, confidence = boxes[sel], class_id[sel], confidence[sel] + sel = nms(boxes, confidence, self.config['iou_threshold']) + boxes, class_id, confidence = boxes[sel], class_id[sel], confidence[sel] + return boxes, class_id, confidence + + def predict(self, payload): + in_path = download_from_youtube(payload['url']) + out_path = f'{uuid.uuid1()}.mp4' + writer = FrameWriter(out_path, size=(416, 416)) + + # TODO Get size from ONNX + for frame in frame_reader(in_path, size=(416, 416)): + x = (frame.astype(np.float32) / 255).transpose(2, 0, 1) + # 4 output tensors, the last three are intermediate values and + # not necessary for detection + output, *_ = self.client.predict(x[None]) + boxes, class_id, confidence = self.postprocess(output) + for x1, y1, x2, y2 in boxes: + frame = cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 1) + writer.write(frame) + + del writer + return FileResponse(out_path) \ No newline at end of file diff --git a/examples/onnx/yolov5-videos/requirements.txt b/examples/onnx/yolov5-videos/requirements.txt new file mode 100644 index 0000000000..2c779ca7f1 --- /dev/null +++ b/examples/onnx/yolov5-videos/requirements.txt @@ -0,0 +1,3 @@ +ffmpeg-python +aiofiles +opencv-python-headless From 313d408416e4e2460247fde0ce400907ce469414 Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Sat, 4 Jul 2020 13:03:03 +1000 Subject: [PATCH 02/11] Update predictor with flexible input size & better overlay --- .../onnx/yolov5-videos/conda-packages.txt | 1 + examples/onnx/yolov5-videos/labels.json | 82 +++++++++++ examples/onnx/yolov5-videos/predictor.py | 134 ++++++++++++------ 3 files changed, 174 insertions(+), 43 deletions(-) create mode 100644 examples/onnx/yolov5-videos/labels.json diff --git a/examples/onnx/yolov5-videos/conda-packages.txt b/examples/onnx/yolov5-videos/conda-packages.txt index 49cbb29eaf..27723c1f77 100644 --- a/examples/onnx/yolov5-videos/conda-packages.txt +++ b/examples/onnx/yolov5-videos/conda-packages.txt @@ -1,2 +1,3 @@ conda-forge::ffmpeg conda-forge::youtube-dl +conda-forge::matplotlib \ No newline at end of file diff --git a/examples/onnx/yolov5-videos/labels.json b/examples/onnx/yolov5-videos/labels.json new file mode 100644 index 0000000000..c4f7d257b1 --- /dev/null +++ b/examples/onnx/yolov5-videos/labels.json @@ -0,0 +1,82 @@ +[ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush" +] \ No newline at end of file diff --git a/examples/onnx/yolov5-videos/predictor.py b/examples/onnx/yolov5-videos/predictor.py index a2289eeea9..aa7b8c9eb7 100644 --- a/examples/onnx/yolov5-videos/predictor.py +++ b/examples/onnx/yolov5-videos/predictor.py @@ -1,23 +1,30 @@ # WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version` -import ffmpeg -import youtube_dl +import json +import os import uuid from pathlib import Path -import numpy as np +from typing import Iterable, Tuple + import cv2 +import ffmpeg +import numpy as np +import youtube_dl +from matplotlib import pyplot as plt + from starlette.responses import FileResponse -from typing import Iterable, Tuple -import os -def download_from_youtube(url: str) -> Path: - target = f'{uuid.uuid1()}.mp4' - ydl_opts = {'outtmpl': target, 'format': "worstvideo[vcodec=vp9][height>=480]"} +def download_from_youtube(url: str, min_height: int) -> Path: + target = f"{uuid.uuid1()}.mp4" + ydl_opts = { + "outtmpl": target, + "format": f"worstvideo[vcodec=vp9][height>={min_height}]", + } with youtube_dl.YoutubeDL(ydl_opts) as ydl: ydl.download([url]) # we need to glob in case youtube-dl adds suffix - path, = Path().absolute().glob(f'{target}*') + (path,) = Path().absolute().glob(f"{target}*") return path @@ -25,12 +32,13 @@ def frame_reader(path: Path, size: Tuple[int, int]) -> Iterable[np.ndarray]: width, height = size # letterbox frames to fixed size process = ( - ffmpeg - .input(path) - .filter('scale', size=f'{width}:{height}', force_original_aspect_ratio='decrease') + ffmpeg.input(path) + .filter( + "scale", size=f"{width}:{height}", force_original_aspect_ratio="decrease" + ) # Negative values for x and y center the padded video - .filter('pad', height=height, width=width, x=-1, y=-1) - .output('pipe:', format='rawvideo', pix_fmt='rgb24') + .filter("pad", height=height, width=width, x=-1, y=-1) + .output("pipe:", format="rawvideo", pix_fmt="rgb24") .run_async(pipe_stdout=True) ) @@ -44,91 +52,131 @@ def frame_reader(path: Path, size: Tuple[int, int]) -> Iterable[np.ndarray]: class FrameWriter: - def __init__(self, path, size): + def __init__(self, path: Path, size: Tuple[int, int]): width, height = size self.process = ( - ffmpeg - .input('pipe:', format='rawvideo', pix_fmt='rgb24', s=f'{width}x{height}') - .output(path, pix_fmt='yuv420p') + ffmpeg.input( + "pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{width}x{height}" + ) + .output(path, pix_fmt="yuv420p") .overwrite_output() .run_async(pipe_stdin=True) ) - def write(self, frame): + def write(self, frame: np.ndarray): self.process.stdin.write(frame.astype(np.uint8).tobytes()) - def __del__(self): self.process.stdin.close() self.process.wait() -def nms(dets, scores, thresh): +def nms(dets: np.ndarray, scores: np.ndarray, thresh: float) -> np.ndarray: x1 = dets[:, 0] y1 = dets[:, 1] x2 = dets[:, 2] y2 = dets[:, 3] areas = (x2 - x1 + 1) * (y2 - y1 + 1) - order = scores.argsort()[::-1] # get boxes with more ious first + order = scores.argsort()[::-1] # get boxes with more ious first keep = [] while order.size > 0: - i = order[0] # pick maxmum iou box + i = order[0] # pick maxmum iou box keep.append(i) xx1 = np.maximum(x1[i], x1[order[1:]]) yy1 = np.maximum(y1[i], y1[order[1:]]) xx2 = np.minimum(x2[i], x2[order[1:]]) yy2 = np.minimum(y2[i], y2[order[1:]]) - w = np.maximum(0.0, xx2 - xx1 + 1) # maximum width - h = np.maximum(0.0, yy2 - yy1 + 1) # maxiumum height + w = np.maximum(0.0, xx2 - xx1 + 1) # maximum width + h = np.maximum(0.0, yy2 - yy1 + 1) # maxiumum height inter = w * h ovr = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(ovr <= thresh)[0] order = order[inds + 1] - return keep - + return np.array(keep).astype(np.int) + + +def boxes_yolo_to_xyxy(boxes: np.ndarray): + boxes[:, 0] -= boxes[:, 2] / 2 + boxes[:, 1] -= boxes[:, 3] / 2 + boxes[:, 2] = boxes[:, 2] + boxes[:, 0] + boxes[:, 3] = boxes[:, 3] + boxes[:, 1] + return boxes + + +def overlay_boxes(frame, boxes, class_ids, label_map, color_map, line_thickness=None): + tl = ( + line_thickness or round(0.002 * (frame.shape[0] + frame.shape[1]) / 2) + 1 + ) # line/font thickness + + for class_id, (x1, y1, x2, y2) in zip(class_ids, boxes.astype(np.int)): + color = color_map[class_id] + label = label_map[class_id] + cv2.rectangle(frame, (x1, y1), (x2, y2), color, tl, cv2.LINE_AA) + tf = max(tl - 1, 1) # font thickness + t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] + x3, y3 = x1 + t_size[0], y1 - t_size[1] - 3 + cv2.rectangle(frame, (x1, y1), (x3, y3), color, -1, cv2.LINE_AA) # filled + cv2.putText( + frame, + label, + (x1, y1 - 2), + 0, + tl / 3, + [225, 255, 255], + thickness=tf, + lineType=cv2.LINE_AA, + ) class ONNXPredictor: def __init__(self, onnx_client, config): self.client = onnx_client + # Get the input shape from the ONNX runtime + (signature,) = onnx_client.input_signatures.values() + _, _, height, width = signature["images"]["shape"] + self.input_size = (width, height) self.config = config + with open("labels.json") as buf: + self.labels = json.load(buf) + color_map = plt.cm.tab20(np.linspace(0, 20, len(self.labels))) + self.color_map = [tuple(map(int, colors)) for colors in 255 * color_map] def postprocess(self, output): boxes, obj_score, class_scores = np.split(output[0], [4, 5], axis=1) - # convert boxes from x_center, y_center, w, h to xyxy - boxes[:, 0] -= boxes[:, 2] / 2 - boxes[:, 1] -= boxes[:, 3] / 2 - boxes[:, 2] = boxes[:, 2] + boxes[:, 0] - boxes[:, 3] = boxes[:, 3] + boxes[:, 1] + boxes = boxes_yolo_to_xyxy(boxes) + + # get the class-prediction & class confidences class_id = class_scores.argmax(axis=1) cls_score = class_scores[np.arange(len(class_scores)), class_id] + confidence = obj_score.squeeze(axis=1) * cls_score - sel = confidence > self.config['confidence_threshold'] + sel = confidence > self.config["confidence_threshold"] boxes, class_id, confidence = boxes[sel], class_id[sel], confidence[sel] - sel = nms(boxes, confidence, self.config['iou_threshold']) + sel = nms(boxes, confidence, self.config["iou_threshold"]) boxes, class_id, confidence = boxes[sel], class_id[sel], confidence[sel] return boxes, class_id, confidence def predict(self, payload): - in_path = download_from_youtube(payload['url']) - out_path = f'{uuid.uuid1()}.mp4' - writer = FrameWriter(out_path, size=(416, 416)) + in_path = download_from_youtube(payload["url"], self.input_size[1]) + out_path = f"{uuid.uuid1()}.mp4" + writer = FrameWriter(out_path, size=self.input_size) - # TODO Get size from ONNX - for frame in frame_reader(in_path, size=(416, 416)): + for frame in frame_reader(in_path, size=self.input_size): x = (frame.astype(np.float32) / 255).transpose(2, 0, 1) # 4 output tensors, the last three are intermediate values and # not necessary for detection output, *_ = self.client.predict(x[None]) - boxes, class_id, confidence = self.postprocess(output) - for x1, y1, x2, y2 in boxes: - frame = cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 1) + boxes, class_ids, confidence = self.postprocess(output) + overlay_boxes(frame, boxes, class_ids, self.labels, self.color_map) writer.write(frame) del writer - return FileResponse(out_path) \ No newline at end of file + os.remove(in_path) + # We cant remove out_path, so the deployment will run out of memory + # sooner or later + return FileResponse(out_path) From 29977d772a623aa148446f8d50ec7db366038d6f Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Sat, 4 Jul 2020 13:21:32 +1000 Subject: [PATCH 03/11] Update readme --- examples/onnx/yolov5-videos/README.md | 58 ++++++++++++++++++++++++- examples/onnx/yolov5-videos/sample.json | 3 ++ 2 files changed, 60 insertions(+), 1 deletion(-) create mode 100644 examples/onnx/yolov5-videos/sample.json diff --git a/examples/onnx/yolov5-videos/README.md b/examples/onnx/yolov5-videos/README.md index a0990367ef..155be31118 100644 --- a/examples/onnx/yolov5-videos/README.md +++ b/examples/onnx/yolov5-videos/README.md @@ -1 +1,57 @@ -TBD +# YoloV5 Detection model + +This example deploys a detection model trained using [ultralytics' yolo repo](https://github.com/ultralytics/yolov5) using ONNX. +We'll use the `yolov5s` model as an example here. +In can be used to run inference on youtube videos and returns the annotated video with bounding boxes. + +The example can be run on both CPU and on GPU hardware. + +## Exporting ONNX + +To export a custom model from the repo, use the [`model/export.py`](https://github.com/ultralytics/yolov5/blob/master/models/export.py) script. +The only change we need to make is to change the line + +``` +model.model[-1].export = True # set Detect() layer export=True +``` + +to + +``` +model.model[-1].export = False +``` + +Originally, the ultralytics repo does not export postprocessing steps of the model, e.g. the conversion from the raw CNN outputs to bounding boxes. +With newer ONNX versions, these can be exported as part of the model making the deployment much easier. + +With this modified script, the ONNX graph used for this example has been exported using +``` +python models/export.py --weights weights/yolov5s.pt --img 416 --batch 1 +``` + + +## Sample Prediction + +Deploy the model by running: + +```bash +cortex deploy +``` + +And wait for it to become live by tracking its status with `cortex get --watch`. + +Once the API has been successfully deployed, export the API's endpoint for convenience. You can get the API's endpoint by running `cortex get youtube-yolov5`. + +```bash +export ENDPOINT=your-api-endpoint +``` + +When making a prediction with [sample.json](sample.json), [this](https://www.youtube.com/watch?v=aUdKzb4LGJ) youtube video will be used. + +To make a request to the model: + +```bash +curl "${ENDPOINT}" -X POST -H "Content-Type: application/json" -d @sample.json --output video.mp4 +``` + +After a few seconds, `curl` will save the resulting video `video.mp4` in the current working directory. \ No newline at end of file diff --git a/examples/onnx/yolov5-videos/sample.json b/examples/onnx/yolov5-videos/sample.json new file mode 100644 index 0000000000..8421278f58 --- /dev/null +++ b/examples/onnx/yolov5-videos/sample.json @@ -0,0 +1,3 @@ +{ + "url": "https://www.youtube.com/watch?v=aUdKzb4LGJI" +} From b67e4ebf9c9c5e276c74d08ea996abcf22945f3f Mon Sep 17 00:00:00 2001 From: Robert Lucian Chiriac Date: Sat, 4 Jul 2020 07:03:48 +0300 Subject: [PATCH 04/11] Make lint --- examples/onnx/yolov5-videos/README.md | 2 +- examples/onnx/yolov5-videos/conda-packages.txt | 2 +- examples/onnx/yolov5-videos/cortex.yaml | 2 +- examples/onnx/yolov5-videos/labels.json | 2 +- examples/onnx/yolov5-videos/predictor.py | 8 ++------ 5 files changed, 6 insertions(+), 10 deletions(-) diff --git a/examples/onnx/yolov5-videos/README.md b/examples/onnx/yolov5-videos/README.md index 155be31118..003c59603b 100644 --- a/examples/onnx/yolov5-videos/README.md +++ b/examples/onnx/yolov5-videos/README.md @@ -54,4 +54,4 @@ To make a request to the model: curl "${ENDPOINT}" -X POST -H "Content-Type: application/json" -d @sample.json --output video.mp4 ``` -After a few seconds, `curl` will save the resulting video `video.mp4` in the current working directory. \ No newline at end of file +After a few seconds, `curl` will save the resulting video `video.mp4` in the current working directory. diff --git a/examples/onnx/yolov5-videos/conda-packages.txt b/examples/onnx/yolov5-videos/conda-packages.txt index 27723c1f77..a3cd6a00e8 100644 --- a/examples/onnx/yolov5-videos/conda-packages.txt +++ b/examples/onnx/yolov5-videos/conda-packages.txt @@ -1,3 +1,3 @@ conda-forge::ffmpeg conda-forge::youtube-dl -conda-forge::matplotlib \ No newline at end of file +conda-forge::matplotlib diff --git a/examples/onnx/yolov5-videos/cortex.yaml b/examples/onnx/yolov5-videos/cortex.yaml index c6c89a344d..f6dfcad3e6 100644 --- a/examples/onnx/yolov5-videos/cortex.yaml +++ b/examples/onnx/yolov5-videos/cortex.yaml @@ -10,4 +10,4 @@ iou_threshold: 0.5 confidence_threshold: 0.3 compute: - gpu: 1 \ No newline at end of file + gpu: 1 diff --git a/examples/onnx/yolov5-videos/labels.json b/examples/onnx/yolov5-videos/labels.json index c4f7d257b1..c86f2f812a 100644 --- a/examples/onnx/yolov5-videos/labels.json +++ b/examples/onnx/yolov5-videos/labels.json @@ -79,4 +79,4 @@ "teddy bear", "hair drier", "toothbrush" -] \ No newline at end of file +] diff --git a/examples/onnx/yolov5-videos/predictor.py b/examples/onnx/yolov5-videos/predictor.py index aa7b8c9eb7..b7b87f2015 100644 --- a/examples/onnx/yolov5-videos/predictor.py +++ b/examples/onnx/yolov5-videos/predictor.py @@ -33,9 +33,7 @@ def frame_reader(path: Path, size: Tuple[int, int]) -> Iterable[np.ndarray]: # letterbox frames to fixed size process = ( ffmpeg.input(path) - .filter( - "scale", size=f"{width}:{height}", force_original_aspect_ratio="decrease" - ) + .filter("scale", size=f"{width}:{height}", force_original_aspect_ratio="decrease") # Negative values for x and y center the padded video .filter("pad", height=height, width=width, x=-1, y=-1) .output("pipe:", format="rawvideo", pix_fmt="rgb24") @@ -55,9 +53,7 @@ class FrameWriter: def __init__(self, path: Path, size: Tuple[int, int]): width, height = size self.process = ( - ffmpeg.input( - "pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{width}x{height}" - ) + ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{width}x{height}") .output(path, pix_fmt="yuv420p") .overwrite_output() .run_async(pipe_stdin=True) From d0caa6c7da8b4f647b2db00a9aeae47c6bb1bffc Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Sun, 5 Jul 2020 21:19:33 +1000 Subject: [PATCH 05/11] Rename to yolov5-youtube --- examples/onnx/{yolov5-videos => yolov5-youtube}/README.md | 0 .../{yolov5-videos => yolov5-youtube}/conda-packages.txt | 0 examples/onnx/{yolov5-videos => yolov5-youtube}/cortex.yaml | 5 +++-- examples/onnx/{yolov5-videos => yolov5-youtube}/labels.json | 0 examples/onnx/{yolov5-videos => yolov5-youtube}/predictor.py | 0 .../onnx/{yolov5-videos => yolov5-youtube}/requirements.txt | 0 examples/onnx/{yolov5-videos => yolov5-youtube}/sample.json | 0 7 files changed, 3 insertions(+), 2 deletions(-) rename examples/onnx/{yolov5-videos => yolov5-youtube}/README.md (100%) rename examples/onnx/{yolov5-videos => yolov5-youtube}/conda-packages.txt (100%) rename examples/onnx/{yolov5-videos => yolov5-youtube}/cortex.yaml (60%) rename examples/onnx/{yolov5-videos => yolov5-youtube}/labels.json (100%) rename examples/onnx/{yolov5-videos => yolov5-youtube}/predictor.py (100%) rename examples/onnx/{yolov5-videos => yolov5-youtube}/requirements.txt (100%) rename examples/onnx/{yolov5-videos => yolov5-youtube}/sample.json (100%) diff --git a/examples/onnx/yolov5-videos/README.md b/examples/onnx/yolov5-youtube/README.md similarity index 100% rename from examples/onnx/yolov5-videos/README.md rename to examples/onnx/yolov5-youtube/README.md diff --git a/examples/onnx/yolov5-videos/conda-packages.txt b/examples/onnx/yolov5-youtube/conda-packages.txt similarity index 100% rename from examples/onnx/yolov5-videos/conda-packages.txt rename to examples/onnx/yolov5-youtube/conda-packages.txt diff --git a/examples/onnx/yolov5-videos/cortex.yaml b/examples/onnx/yolov5-youtube/cortex.yaml similarity index 60% rename from examples/onnx/yolov5-videos/cortex.yaml rename to examples/onnx/yolov5-youtube/cortex.yaml index f6dfcad3e6..81619a88a1 100644 --- a/examples/onnx/yolov5-videos/cortex.yaml +++ b/examples/onnx/yolov5-youtube/cortex.yaml @@ -1,13 +1,14 @@ # WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version` -- name: youtube-yolov5 +- name: yolov5-youtube kind: SyncAPI predictor: type: onnx path: predictor.py - model_path: yolov5s.onnx + model_path: s3://cortex-examples/onnx/yolov5-youtube/yolov5s.onnx config: iou_threshold: 0.5 confidence_threshold: 0.3 compute: + # GPU requirement is optional. Comment out next line to run on CPUs (albeit slower) gpu: 1 diff --git a/examples/onnx/yolov5-videos/labels.json b/examples/onnx/yolov5-youtube/labels.json similarity index 100% rename from examples/onnx/yolov5-videos/labels.json rename to examples/onnx/yolov5-youtube/labels.json diff --git a/examples/onnx/yolov5-videos/predictor.py b/examples/onnx/yolov5-youtube/predictor.py similarity index 100% rename from examples/onnx/yolov5-videos/predictor.py rename to examples/onnx/yolov5-youtube/predictor.py diff --git a/examples/onnx/yolov5-videos/requirements.txt b/examples/onnx/yolov5-youtube/requirements.txt similarity index 100% rename from examples/onnx/yolov5-videos/requirements.txt rename to examples/onnx/yolov5-youtube/requirements.txt diff --git a/examples/onnx/yolov5-videos/sample.json b/examples/onnx/yolov5-youtube/sample.json similarity index 100% rename from examples/onnx/yolov5-videos/sample.json rename to examples/onnx/yolov5-youtube/sample.json From 9de3281fefd7339b5b1f6ef990b515f33d9d9850 Mon Sep 17 00:00:00 2001 From: Daniel Suess Date: Sun, 5 Jul 2020 21:22:11 +1000 Subject: [PATCH 06/11] Update Readme --- examples/onnx/yolov5-youtube/README.md | 50 +++++++++++------------ examples/onnx/yolov5-youtube/predictor.py | 27 ++++++------ 2 files changed, 40 insertions(+), 37 deletions(-) diff --git a/examples/onnx/yolov5-youtube/README.md b/examples/onnx/yolov5-youtube/README.md index 003c59603b..c8d8d016a3 100644 --- a/examples/onnx/yolov5-youtube/README.md +++ b/examples/onnx/yolov5-youtube/README.md @@ -6,30 +6,6 @@ In can be used to run inference on youtube videos and returns the annotated vide The example can be run on both CPU and on GPU hardware. -## Exporting ONNX - -To export a custom model from the repo, use the [`model/export.py`](https://github.com/ultralytics/yolov5/blob/master/models/export.py) script. -The only change we need to make is to change the line - -``` -model.model[-1].export = True # set Detect() layer export=True -``` - -to - -``` -model.model[-1].export = False -``` - -Originally, the ultralytics repo does not export postprocessing steps of the model, e.g. the conversion from the raw CNN outputs to bounding boxes. -With newer ONNX versions, these can be exported as part of the model making the deployment much easier. - -With this modified script, the ONNX graph used for this example has been exported using -``` -python models/export.py --weights weights/yolov5s.pt --img 416 --batch 1 -``` - - ## Sample Prediction Deploy the model by running: @@ -40,7 +16,7 @@ cortex deploy And wait for it to become live by tracking its status with `cortex get --watch`. -Once the API has been successfully deployed, export the API's endpoint for convenience. You can get the API's endpoint by running `cortex get youtube-yolov5`. +Once the API has been successfully deployed, export the API's endpoint for convenience. You can get the API's endpoint by running `cortex get yolov5-youtube`. ```bash export ENDPOINT=your-api-endpoint @@ -55,3 +31,27 @@ curl "${ENDPOINT}" -X POST -H "Content-Type: application/json" -d @sample.json - ``` After a few seconds, `curl` will save the resulting video `video.mp4` in the current working directory. + + +## Exporting ONNX + +To export a custom model from the repo, use the [`model/export.py`](https://github.com/ultralytics/yolov5/blob/master/models/export.py) script. +The only change we need to make is to change the line + +```bash +model.model[-1].export = True # set Detect() layer export=True +``` + +to + +```bash +model.model[-1].export = False +``` + +Originally, the ultralytics repo does not export postprocessing steps of the model, e.g. the conversion from the raw CNN outputs to bounding boxes. +With newer ONNX versions, these can be exported as part of the model making the deployment much easier. + +With this modified script, the ONNX graph used for this example has been exported using +```bash +python models/export.py --weights weights/yolov5s.pt --img 416 --batch 1 +``` diff --git a/examples/onnx/yolov5-youtube/predictor.py b/examples/onnx/yolov5-youtube/predictor.py index b7b87f2015..c525de2675 100644 --- a/examples/onnx/yolov5-youtube/predictor.py +++ b/examples/onnx/yolov5-youtube/predictor.py @@ -162,17 +162,20 @@ def predict(self, payload): out_path = f"{uuid.uuid1()}.mp4" writer = FrameWriter(out_path, size=self.input_size) - for frame in frame_reader(in_path, size=self.input_size): - x = (frame.astype(np.float32) / 255).transpose(2, 0, 1) - # 4 output tensors, the last three are intermediate values and - # not necessary for detection - output, *_ = self.client.predict(x[None]) - boxes, class_ids, confidence = self.postprocess(output) - overlay_boxes(frame, boxes, class_ids, self.labels, self.color_map) - writer.write(frame) - - del writer - os.remove(in_path) - # We cant remove out_path, so the deployment will run out of memory + try: + for frame in frame_reader(in_path, size=self.input_size): + x = (frame.astype(np.float32) / 255).transpose(2, 0, 1) + # 4 output tensors, the last three are intermediate values and + # not necessary for detection + output, *_ = self.client.predict(x[None]) + boxes, class_ids, confidence = self.postprocess(output) + overlay_boxes(frame, boxes, class_ids, self.labels, self.color_map) + writer.write(frame) + finally: + # We delete the writer manually to close the opened file + del writer + os.remove(in_path) + + # We cant remove out_path, so the deployment will run out of diskmemory # sooner or later return FileResponse(out_path) From d81b0ecbad8b947354e668b76d7a1097b12f96a9 Mon Sep 17 00:00:00 2001 From: Robert Lucian Chiriac Date: Sun, 5 Jul 2020 21:26:00 +0300 Subject: [PATCH 07/11] Fix ffmpeg bug with gmp dependency Fixes error "conda libgnutls.so.30: symbol mpn_add_1 version HOGWEED_4 not defined in file libhogweed.so.4 with link time reference" --- examples/onnx/yolov5-youtube/conda-packages.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/onnx/yolov5-youtube/conda-packages.txt b/examples/onnx/yolov5-youtube/conda-packages.txt index a3cd6a00e8..131fce12b5 100644 --- a/examples/onnx/yolov5-youtube/conda-packages.txt +++ b/examples/onnx/yolov5-youtube/conda-packages.txt @@ -1,3 +1,3 @@ -conda-forge::ffmpeg +conda-forge::ffmpeg=4.2.3 conda-forge::youtube-dl conda-forge::matplotlib From 4621c290406b857fb5ed314c9e98ee758511a92d Mon Sep 17 00:00:00 2001 From: Robert Lucian Chiriac Date: Mon, 6 Jul 2020 01:32:42 +0300 Subject: [PATCH 08/11] Use context manager & remove all videos from disk --- examples/onnx/yolov5-youtube/predictor.py | 25 ++++++++++++++--------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/examples/onnx/yolov5-youtube/predictor.py b/examples/onnx/yolov5-youtube/predictor.py index c525de2675..f2a00a8c30 100644 --- a/examples/onnx/yolov5-youtube/predictor.py +++ b/examples/onnx/yolov5-youtube/predictor.py @@ -62,6 +62,12 @@ def __init__(self, path: Path, size: Tuple[int, int]): def write(self, frame: np.ndarray): self.process.stdin.write(frame.astype(np.uint8).tobytes()) + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.__del__() + def __del__(self): self.process.stdin.close() self.process.wait() @@ -160,9 +166,8 @@ def postprocess(self, output): def predict(self, payload): in_path = download_from_youtube(payload["url"], self.input_size[1]) out_path = f"{uuid.uuid1()}.mp4" - writer = FrameWriter(out_path, size=self.input_size) - try: + with FrameWriter(out_path, size=self.input_size) as writer: for frame in frame_reader(in_path, size=self.input_size): x = (frame.astype(np.float32) / 255).transpose(2, 0, 1) # 4 output tensors, the last three are intermediate values and @@ -171,11 +176,11 @@ def predict(self, payload): boxes, class_ids, confidence = self.postprocess(output) overlay_boxes(frame, boxes, class_ids, self.labels, self.color_map) writer.write(frame) - finally: - # We delete the writer manually to close the opened file - del writer - os.remove(in_path) - - # We cant remove out_path, so the deployment will run out of diskmemory - # sooner or later - return FileResponse(out_path) + + with open(out_path, "rb") as f: + output_bytes = f.read() + + os.remove(in_path) + os.remove(out_path) + + return output_bytes From 8d65d4b3eb27aa4fb2a9e8949c1a45dc716050e6 Mon Sep 17 00:00:00 2001 From: Robert Lucian Chiriac Date: Mon, 6 Jul 2020 01:56:34 +0300 Subject: [PATCH 09/11] Reorganize the example a bit --- examples/onnx/yolov5-youtube/predictor.py | 142 ++-------------------- examples/onnx/yolov5-youtube/utils.py | 130 ++++++++++++++++++++ 2 files changed, 139 insertions(+), 133 deletions(-) create mode 100644 examples/onnx/yolov5-youtube/utils.py diff --git a/examples/onnx/yolov5-youtube/predictor.py b/examples/onnx/yolov5-youtube/predictor.py index f2a00a8c30..2dca90a605 100644 --- a/examples/onnx/yolov5-youtube/predictor.py +++ b/examples/onnx/yolov5-youtube/predictor.py @@ -3,137 +3,11 @@ import json import os import uuid -from pathlib import Path -from typing import Iterable, Tuple +import utils -import cv2 -import ffmpeg import numpy as np -import youtube_dl from matplotlib import pyplot as plt -from starlette.responses import FileResponse - - -def download_from_youtube(url: str, min_height: int) -> Path: - target = f"{uuid.uuid1()}.mp4" - ydl_opts = { - "outtmpl": target, - "format": f"worstvideo[vcodec=vp9][height>={min_height}]", - } - with youtube_dl.YoutubeDL(ydl_opts) as ydl: - ydl.download([url]) - # we need to glob in case youtube-dl adds suffix - (path,) = Path().absolute().glob(f"{target}*") - return path - - -def frame_reader(path: Path, size: Tuple[int, int]) -> Iterable[np.ndarray]: - width, height = size - # letterbox frames to fixed size - process = ( - ffmpeg.input(path) - .filter("scale", size=f"{width}:{height}", force_original_aspect_ratio="decrease") - # Negative values for x and y center the padded video - .filter("pad", height=height, width=width, x=-1, y=-1) - .output("pipe:", format="rawvideo", pix_fmt="rgb24") - .run_async(pipe_stdout=True) - ) - - while True: - in_bytes = process.stdout.read(height * width * 3) - if not in_bytes: - process.wait() - break - frame = np.frombuffer(in_bytes, np.uint8).reshape([height, width, 3]) - yield frame - - -class FrameWriter: - def __init__(self, path: Path, size: Tuple[int, int]): - width, height = size - self.process = ( - ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{width}x{height}") - .output(path, pix_fmt="yuv420p") - .overwrite_output() - .run_async(pipe_stdin=True) - ) - - def write(self, frame: np.ndarray): - self.process.stdin.write(frame.astype(np.uint8).tobytes()) - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, traceback): - self.__del__() - - def __del__(self): - self.process.stdin.close() - self.process.wait() - - -def nms(dets: np.ndarray, scores: np.ndarray, thresh: float) -> np.ndarray: - x1 = dets[:, 0] - y1 = dets[:, 1] - x2 = dets[:, 2] - y2 = dets[:, 3] - - areas = (x2 - x1 + 1) * (y2 - y1 + 1) - order = scores.argsort()[::-1] # get boxes with more ious first - - keep = [] - while order.size > 0: - i = order[0] # pick maxmum iou box - keep.append(i) - xx1 = np.maximum(x1[i], x1[order[1:]]) - yy1 = np.maximum(y1[i], y1[order[1:]]) - xx2 = np.minimum(x2[i], x2[order[1:]]) - yy2 = np.minimum(y2[i], y2[order[1:]]) - - w = np.maximum(0.0, xx2 - xx1 + 1) # maximum width - h = np.maximum(0.0, yy2 - yy1 + 1) # maxiumum height - inter = w * h - ovr = inter / (areas[i] + areas[order[1:]] - inter) - - inds = np.where(ovr <= thresh)[0] - order = order[inds + 1] - - return np.array(keep).astype(np.int) - - -def boxes_yolo_to_xyxy(boxes: np.ndarray): - boxes[:, 0] -= boxes[:, 2] / 2 - boxes[:, 1] -= boxes[:, 3] / 2 - boxes[:, 2] = boxes[:, 2] + boxes[:, 0] - boxes[:, 3] = boxes[:, 3] + boxes[:, 1] - return boxes - - -def overlay_boxes(frame, boxes, class_ids, label_map, color_map, line_thickness=None): - tl = ( - line_thickness or round(0.002 * (frame.shape[0] + frame.shape[1]) / 2) + 1 - ) # line/font thickness - - for class_id, (x1, y1, x2, y2) in zip(class_ids, boxes.astype(np.int)): - color = color_map[class_id] - label = label_map[class_id] - cv2.rectangle(frame, (x1, y1), (x2, y2), color, tl, cv2.LINE_AA) - tf = max(tl - 1, 1) # font thickness - t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] - x3, y3 = x1 + t_size[0], y1 - t_size[1] - 3 - cv2.rectangle(frame, (x1, y1), (x3, y3), color, -1, cv2.LINE_AA) # filled - cv2.putText( - frame, - label, - (x1, y1 - 2), - 0, - tl / 3, - [225, 255, 255], - thickness=tf, - lineType=cv2.LINE_AA, - ) - class ONNXPredictor: def __init__(self, onnx_client, config): @@ -150,7 +24,7 @@ def __init__(self, onnx_client, config): def postprocess(self, output): boxes, obj_score, class_scores = np.split(output[0], [4, 5], axis=1) - boxes = boxes_yolo_to_xyxy(boxes) + boxes = utils.boxes_yolo_to_xyxy(boxes) # get the class-prediction & class confidences class_id = class_scores.argmax(axis=1) @@ -159,22 +33,24 @@ def postprocess(self, output): confidence = obj_score.squeeze(axis=1) * cls_score sel = confidence > self.config["confidence_threshold"] boxes, class_id, confidence = boxes[sel], class_id[sel], confidence[sel] - sel = nms(boxes, confidence, self.config["iou_threshold"]) + sel = utils.nms(boxes, confidence, self.config["iou_threshold"]) boxes, class_id, confidence = boxes[sel], class_id[sel], confidence[sel] return boxes, class_id, confidence def predict(self, payload): - in_path = download_from_youtube(payload["url"], self.input_size[1]) + # download YT video + in_path = utils.download_from_youtube(payload["url"], self.input_size[1]) out_path = f"{uuid.uuid1()}.mp4" - with FrameWriter(out_path, size=self.input_size) as writer: - for frame in frame_reader(in_path, size=self.input_size): + # run predictions + with utils.FrameWriter(out_path, size=self.input_size) as writer: + for frame in utils.frame_reader(in_path, size=self.input_size): x = (frame.astype(np.float32) / 255).transpose(2, 0, 1) # 4 output tensors, the last three are intermediate values and # not necessary for detection output, *_ = self.client.predict(x[None]) boxes, class_ids, confidence = self.postprocess(output) - overlay_boxes(frame, boxes, class_ids, self.labels, self.color_map) + utils.overlay_boxes(frame, boxes, class_ids, self.labels, self.color_map) writer.write(frame) with open(out_path, "rb") as f: diff --git a/examples/onnx/yolov5-youtube/utils.py b/examples/onnx/yolov5-youtube/utils.py new file mode 100644 index 0000000000..6ccf9e7871 --- /dev/null +++ b/examples/onnx/yolov5-youtube/utils.py @@ -0,0 +1,130 @@ +# WARNING: you are on the master branch, please refer to the examples on the branch that matches your `cortex version` + +import youtube_dl +import ffmpeg +import numpy as np +import cv2 +import uuid + +from pathlib import Path +from typing import Iterable, Tuple + + +def download_from_youtube(url: str, min_height: int) -> Path: + target = f"{uuid.uuid1()}.mp4" + ydl_opts = { + "outtmpl": target, + "format": f"worstvideo[vcodec=vp9][height>={min_height}]", + } + with youtube_dl.YoutubeDL(ydl_opts) as ydl: + ydl.download([url]) + # we need to glob in case youtube-dl adds suffix + (path,) = Path().absolute().glob(f"{target}*") + return path + + +def frame_reader(path: Path, size: Tuple[int, int]) -> Iterable[np.ndarray]: + width, height = size + # letterbox frames to fixed size + process = ( + ffmpeg.input(path) + .filter("scale", size=f"{width}:{height}", force_original_aspect_ratio="decrease") + # Negative values for x and y center the padded video + .filter("pad", height=height, width=width, x=-1, y=-1) + .output("pipe:", format="rawvideo", pix_fmt="rgb24") + .run_async(pipe_stdout=True) + ) + + while True: + in_bytes = process.stdout.read(height * width * 3) + if not in_bytes: + process.wait() + break + frame = np.frombuffer(in_bytes, np.uint8).reshape([height, width, 3]) + yield frame + + +class FrameWriter: + def __init__(self, path: Path, size: Tuple[int, int]): + width, height = size + self.process = ( + ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", s=f"{width}x{height}") + .output(path, pix_fmt="yuv420p") + .overwrite_output() + .run_async(pipe_stdin=True) + ) + + def write(self, frame: np.ndarray): + self.process.stdin.write(frame.astype(np.uint8).tobytes()) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.__del__() + + def __del__(self): + self.process.stdin.close() + self.process.wait() + + +def nms(dets: np.ndarray, scores: np.ndarray, thresh: float) -> np.ndarray: + x1 = dets[:, 0] + y1 = dets[:, 1] + x2 = dets[:, 2] + y2 = dets[:, 3] + + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] # get boxes with more ious first + + keep = [] + while order.size > 0: + i = order[0] # pick maxmum iou box + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) # maximum width + h = np.maximum(0.0, yy2 - yy1 + 1) # maxiumum height + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + inds = np.where(ovr <= thresh)[0] + order = order[inds + 1] + + return np.array(keep).astype(np.int) + + +def boxes_yolo_to_xyxy(boxes: np.ndarray): + boxes[:, 0] -= boxes[:, 2] / 2 + boxes[:, 1] -= boxes[:, 3] / 2 + boxes[:, 2] = boxes[:, 2] + boxes[:, 0] + boxes[:, 3] = boxes[:, 3] + boxes[:, 1] + return boxes + + +def overlay_boxes(frame, boxes, class_ids, label_map, color_map, line_thickness=None): + tl = ( + line_thickness or round(0.002 * (frame.shape[0] + frame.shape[1]) / 2) + 1 + ) # line/font thickness + + for class_id, (x1, y1, x2, y2) in zip(class_ids, boxes.astype(np.int)): + color = color_map[class_id] + label = label_map[class_id] + cv2.rectangle(frame, (x1, y1), (x2, y2), color, tl, cv2.LINE_AA) + tf = max(tl - 1, 1) # font thickness + t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] + x3, y3 = x1 + t_size[0], y1 - t_size[1] - 3 + cv2.rectangle(frame, (x1, y1), (x3, y3), color, -1, cv2.LINE_AA) # filled + cv2.putText( + frame, + label, + (x1, y1 - 2), + 0, + tl / 3, + [225, 255, 255], + thickness=tf, + lineType=cv2.LINE_AA, + ) From 436bcd332b3c2817dc90915b499645099723f984 Mon Sep 17 00:00:00 2001 From: Robert Lucian Chiriac Date: Mon, 6 Jul 2020 03:16:51 +0300 Subject: [PATCH 10/11] Polishing the docs & tuning line thickness --- examples/README.md | 2 ++ examples/onnx/yolov5-youtube/README.md | 6 ++++-- examples/onnx/yolov5-youtube/cortex.yaml | 2 +- examples/onnx/yolov5-youtube/utils.py | 2 +- 4 files changed, 8 insertions(+), 4 deletions(-) diff --git a/examples/README.md b/examples/README.md index 5ae53860c8..21e585caf6 100644 --- a/examples/README.md +++ b/examples/README.md @@ -50,6 +50,8 @@ - [Iris classification](onnx/iris-classifier): deploy an XGBoost model (exported in ONNX) to classify iris flowers. +- [YOLOv5 YouTube detection](onnx/yolov5-youtube): deploy a YOLOv5 model trained on COCO val2017 dataset. + - [Multi-model classification](onnx/multi-model-classifier): deploy 3 models (ResNet50, MobileNet, ShuffleNet) in a single API. ## scikit-learn diff --git a/examples/onnx/yolov5-youtube/README.md b/examples/onnx/yolov5-youtube/README.md index c8d8d016a3..2b7d8a0deb 100644 --- a/examples/onnx/yolov5-youtube/README.md +++ b/examples/onnx/yolov5-youtube/README.md @@ -1,4 +1,4 @@ -# YoloV5 Detection model +# YOLOv5 Detection model This example deploys a detection model trained using [ultralytics' yolo repo](https://github.com/ultralytics/yolov5) using ONNX. We'll use the `yolov5s` model as an example here. @@ -30,7 +30,9 @@ To make a request to the model: curl "${ENDPOINT}" -X POST -H "Content-Type: application/json" -d @sample.json --output video.mp4 ``` -After a few seconds, `curl` will save the resulting video `video.mp4` in the current working directory. +After a few seconds, `curl` will save the resulting video `video.mp4` in the current working directory. The following is a sample of what should be exported: + +![yolov5](https://user-images.githubusercontent.com/26958764/86545098-e0dce900-bf34-11ea-83a7-8fd544afa11c.gif) ## Exporting ONNX diff --git a/examples/onnx/yolov5-youtube/cortex.yaml b/examples/onnx/yolov5-youtube/cortex.yaml index 81619a88a1..0dcbc231d1 100644 --- a/examples/onnx/yolov5-youtube/cortex.yaml +++ b/examples/onnx/yolov5-youtube/cortex.yaml @@ -8,7 +8,7 @@ model_path: s3://cortex-examples/onnx/yolov5-youtube/yolov5s.onnx config: iou_threshold: 0.5 - confidence_threshold: 0.3 + confidence_threshold: 0.6 compute: # GPU requirement is optional. Comment out next line to run on CPUs (albeit slower) gpu: 1 diff --git a/examples/onnx/yolov5-youtube/utils.py b/examples/onnx/yolov5-youtube/utils.py index 6ccf9e7871..3e8cfff9c8 100644 --- a/examples/onnx/yolov5-youtube/utils.py +++ b/examples/onnx/yolov5-youtube/utils.py @@ -107,7 +107,7 @@ def boxes_yolo_to_xyxy(boxes: np.ndarray): def overlay_boxes(frame, boxes, class_ids, label_map, color_map, line_thickness=None): tl = ( - line_thickness or round(0.002 * (frame.shape[0] + frame.shape[1]) / 2) + 1 + line_thickness or round(0.0005 * (frame.shape[0] + frame.shape[1]) / 2) + 1 ) # line/font thickness for class_id, (x1, y1, x2, y2) in zip(class_ids, boxes.astype(np.int)): From 001f40f1a03e7a7b9424df07aec4fb8940e1d3c3 Mon Sep 17 00:00:00 2001 From: Robert Lucian Chiriac Date: Mon, 6 Jul 2020 04:19:05 +0300 Subject: [PATCH 11/11] Use video/mp4 mime-type --- examples/onnx/yolov5-youtube/predictor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/onnx/yolov5-youtube/predictor.py b/examples/onnx/yolov5-youtube/predictor.py index 2dca90a605..be17e13139 100644 --- a/examples/onnx/yolov5-youtube/predictor.py +++ b/examples/onnx/yolov5-youtube/predictor.py @@ -2,12 +2,15 @@ import json import os +import io import uuid import utils import numpy as np from matplotlib import pyplot as plt +from starlette.responses import StreamingResponse + class ONNXPredictor: def __init__(self, onnx_client, config): @@ -54,9 +57,9 @@ def predict(self, payload): writer.write(frame) with open(out_path, "rb") as f: - output_bytes = f.read() + output_buf = io.BytesIO(f.read()) os.remove(in_path) os.remove(out_path) - return output_bytes + return StreamingResponse(output_buf, media_type="video/mp4")