From 966b0e09f0a5261e555c2a137af2ef9d58cc9779 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 25 Sep 2022 16:21:26 +0200 Subject: [PATCH] TensorRT detect.py inference fix (#9581) * Update * Update ci-testing.yml Signed-off-by: Glenn Jocher * Update ci-testing.yml Signed-off-by: Glenn Jocher * Segment fix * Segment fix Signed-off-by: Glenn Jocher --- .github/workflows/ci-testing.yml | 6 ++++++ classify/predict.py | 3 ++- detect.py | 3 ++- segment/predict.py | 5 +++-- 4 files changed, 13 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci-testing.yml b/.github/workflows/ci-testing.yml index fffc92d1b72f..1ec68e8412f9 100644 --- a/.github/workflows/ci-testing.yml +++ b/.github/workflows/ci-testing.yml @@ -44,6 +44,12 @@ jobs: - name: Benchmark SegmentationModel run: | python benchmarks.py --data coco128-seg.yaml --weights ${{ matrix.model }}-seg.pt --img 320 --hard-fail 0.22 + - name: Test predictions + run: | + python export.py --weights ${{ matrix.model }}-cls.pt --include onnx --img 224 + python detect.py --weights ${{ matrix.model }}.onnx --img 320 + python segment/predict.py --weights ${{ matrix.model }}-seg.onnx --img 320 + python classify/predict.py --weights ${{ matrix.model }}-cls.onnx --img 224 Tests: timeout-minutes: 60 diff --git a/classify/predict.py b/classify/predict.py index d3bec8eea7ba..9114aab1d703 100644 --- a/classify/predict.py +++ b/classify/predict.py @@ -89,14 +89,15 @@ def run( imgsz = check_img_size(imgsz, s=stride) # check image size # Dataloader + bs = 1 # batch_size if webcam: view_img = check_imshow() dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride) + bs = len(dataset) elif screenshot: dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) else: dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride) - bs = len(dataset) # batch_size vid_path, vid_writer = [None] * bs, [None] * bs # Run inference diff --git a/detect.py b/detect.py index 4971033b35fb..8f48d8d28000 100644 --- a/detect.py +++ b/detect.py @@ -97,14 +97,15 @@ def run( imgsz = check_img_size(imgsz, s=stride) # check image size # Dataloader + bs = 1 # batch_size if webcam: view_img = check_imshow() dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) + bs = len(dataset) elif screenshot: dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) else: dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) - bs = len(dataset) # batch_size vid_path, vid_writer = [None] * bs, [None] * bs # Run inference diff --git a/segment/predict.py b/segment/predict.py index 607a8697d731..94117cd78633 100644 --- a/segment/predict.py +++ b/segment/predict.py @@ -100,14 +100,15 @@ def run( imgsz = check_img_size(imgsz, s=stride) # check image size # Dataloader + bs = 1 # batch_size if webcam: view_img = check_imshow() dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) + bs = len(dataset) elif screenshot: dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) else: dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) - bs = len(dataset) # batch_size vid_path, vid_writer = [None] * bs, [None] * bs # Run inference @@ -179,7 +180,7 @@ def run( c = int(cls) # integer class label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}') annotator.box_label(xyxy, label, color=colors(c, True)) - annotator.draw.polygon(segments[j], outline=colors(c, True), width=3) + # annotator.draw.polygon(segments[j], outline=colors(c, True), width=3) if save_crop: save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)