TensorRT detect.py inference fix (ultralytics#9581)

* Update * Update ci-testing.yml Signed-off-by: Glenn Jocher <[email protected]> * Update ci-testing.yml Signed-off-by: Glenn Jocher <[email protected]> * Segment fix * Segment fix Signed-off-by: Glenn Jocher <[email protected]>
libaos · Sep 25, 2022 · 966b0e0 · 966b0e0
1 parent 2787ad7
commit 966b0e0
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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)