From b25d5a75f2c89aace5cae342f3fe29dfdd46e401 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 22 Sep 2022 23:23:40 +0200 Subject: [PATCH] Refactor dataset batch-size (#9551) --- classify/predict.py | 3 +-- detect.py | 3 +-- segment/predict.py | 3 +-- 3 files changed, 3 insertions(+), 6 deletions(-) diff --git a/classify/predict.py b/classify/predict.py index 4857c69..ef59ff6 100644 --- 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 index 310d169..4015b9a 100644 --- 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 index ba4cf29..2ea6bd9 100644 --- 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