Refactor dataset batch-size (ultralytics#9551)

libaos · Sep 22, 2022 · b25d5a7 · b25d5a7
1 parent 489920a
commit b25d5a7
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/classify/predict.py b/classify/predict.py
@@ -91,10 +91,9 @@ def run(
     if webcam:
         view_img = check_imshow()
         dataset = LoadStreams(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
-        bs = len(dataset)  # batch_size
     else:
         dataset = LoadImages(source, img_size=imgsz, transforms=classify_transforms(imgsz[0]), vid_stride=vid_stride)
-        bs = 1  # batch_size
+    bs = len(dataset)  # batch_size
     vid_path, vid_writer = [None] * bs, [None] * bs
 
     # Run inference

diff --git a/detect.py b/detect.py
@@ -99,10 +99,9 @@ def run(
     if webcam:
         view_img = check_imshow()
         dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
-        bs = len(dataset)  # batch_size
     else:
         dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
-        bs = 1  # batch_size
+    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
@@ -101,10 +101,9 @@ def run(
     if webcam:
         view_img = check_imshow()
         dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
-        bs = len(dataset)  # batch_size
     else:
         dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
-        bs = 1  # batch_size
+    bs = len(dataset)  # batch_size
     vid_path, vid_writer = [None] * bs, [None] * bs
 
     # Run inference