improved inference

dataloader.py

@@ -116,31 +116,12 @@ class YOLOFrameDataset(Dataset):
         self.original_shape = (self.ydim, self.xdim)
         self.shape = np.ceil(np.array(shape) * img_size / stride + pad).astype(int) * stride
 
-        self.batches = []
+        self.batch_indices = []
         for i in range(0,n,batch_size): 
-            self.batches.append((i, min(n, i+batch_size)))
-    
-    @classmethod
-    def load_image(cls, img, img_size=896):
-        """Loads and resizes 1 image from dataset, returns img, original hw, resized hw.
-        Modified from ScaledYOLOv4.datasets.load_image()
-        """
-
-        h0, w0 = img.shape[:2]
-        h1, w1 = h0, w0
-        r = img_size / max(h0, w0)
-        if r != 1:  # always resize down, only resize up if training with augmentation
-            interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
-            img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
-            h1, w1 = img.shape[:2]
-        
-        return img, (h0, w0), (h1, w1)  # img, hw_original, hw_resized
+            self.batch_indices.append((i, min(n, i+batch_size)))
 
-    def __len__(self):
-        return len(self.batches)
-
-    def __iter__(self):
-        for batch_idx in self.batches:
+        self.batches = []
+        for batch_idx in self.batch_indices:
 
             batch = []
             labels = None
@@ -164,7 +145,30 @@ class YOLOFrameDataset(Dataset):
             
             image = torch.stack(batch)
 
-            yield (image, labels, shapes)
+            self.batches.append((image, labels, shapes))
+    
+    @classmethod
+    def load_image(cls, img, img_size=896):
+        """Loads and resizes 1 image from dataset, returns img, original hw, resized hw.
+        Modified from ScaledYOLOv4.datasets.load_image()
+        """
+
+        h0, w0 = img.shape[:2]
+        h1, w1 = h0, w0
+        r = img_size / max(h0, w0)
+        if r != 1:  # always resize down, only resize up if training with augmentation
+            interp = cv2.INTER_AREA if r < 1 else cv2.INTER_LINEAR
+            img = cv2.resize(img, (int(w0 * r), int(h0 * r)), interpolation=interp)
+            h1, w1 = img.shape[:2]
+        
+        return img, (h0, w0), (h1, w1)  # img, hw_original, hw_resized
+
+    def __len__(self):
+        return len(self.batches)
+
+    def __iter__(self):
+        for batch in self.batches:
+            yield batch
 
 class ARISBatchedDataset(Dataset):
     def __init__(self, aris_filepath, beam_width_dir, annotations_file, batch_size, num_frames_bg_subtract=1000, disable_output=False,

inference.py

@@ -125,7 +125,6 @@ def do_detection(dataloader, model, device, gp=None, batch_size=BATCH_SIZE, verb
             size = tuple(img.shape)
             nb, _, height, width = size  # batch size, channels, height, width
 
-            print(nb, _, height, width)
             # Run model & NMS
             with torch.no_grad():
                 inf_out, _ = model(img, augment=False) 

scripts/infer_frames.py

@@ -39,7 +39,7 @@ def main(args, config={}, verbose=True):
 
     dirname = args.frames
     
-    locations = ["test"]
+    locations = ["kenai-val"]
     for loc in locations:
 
         in_loc_dir = os.path.join(dirname, loc)