From b8557f87e3e983ef541fa34a245a16fc30e98455 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 24 Jun 2020 13:02:27 -0700
Subject: [PATCH] add stride to datasets.py

---
 test.py           | 1 +
 train.py          | 6 ++++--
 utils/datasets.py | 4 ++--
 3 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/test.py b/test.py
index b3ff968..ae6451f 100644
--- a/test.py
+++ b/test.py
@@ -73,6 +73,7 @@ def test(data,
                                       batch_size,
                                       rect=True,  # rectangular inference
                                       single_cls=opt.single_cls,  # single class mode
+                                      stride=int(max(model.stride)),  # model stride
                                       pad=0.5)  # padding
         batch_size = min(batch_size, len(dataset))
         nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
diff --git a/train.py b/train.py
index 2dfe781..d9d2884 100644
--- a/train.py
+++ b/train.py
@@ -160,7 +160,8 @@ def train(hyp):
                                   hyp=hyp,  # augmentation hyperparameters
                                   rect=opt.rect,  # rectangular training
                                   cache_images=opt.cache_images,
-                                  single_cls=opt.single_cls)
+                                  single_cls=opt.single_cls,
+                                  stride=gs)
     mlc = np.concatenate(dataset.labels, 0)[:, 0].max()  # max label class
     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Correct your labels or your model.' % (mlc, nc, opt.cfg)
 
@@ -179,7 +180,8 @@ def train(hyp):
                                                                  hyp=hyp,
                                                                  rect=True,
                                                                  cache_images=opt.cache_images,
-                                                                 single_cls=opt.single_cls),
+                                                                 single_cls=opt.single_cls,
+                                                                 stride=gs),
                                              batch_size=batch_size,
                                              num_workers=nw,
                                              pin_memory=True,
diff --git a/utils/datasets.py b/utils/datasets.py
index 65e642d..37d773e 100755
--- a/utils/datasets.py
+++ b/utils/datasets.py
@@ -258,7 +258,7 @@ class LoadStreams:  # multiple IP or RTSP cameras
 
 class LoadImagesAndLabels(Dataset):  # for training/testing
     def __init__(self, path, img_size=416, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
-                 cache_images=False, single_cls=False, pad=0.0):
+                 cache_images=False, single_cls=False, stride=32, pad=0.0):
         try:
             path = str(Path(path))  # os-agnostic
             parent = str(Path(path).parent) + os.sep
@@ -325,7 +325,7 @@ class LoadImagesAndLabels(Dataset):  # for training/testing
                 elif mini > 1:
                     shapes[i] = [1, 1 / mini]
 
-            self.batch_shapes = np.ceil(np.array(shapes) * img_size / 32. + pad).astype(np.int) * 32
+            self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride
 
         # Cache labels
         self.imgs = [None] * n