fixed padding

dataloader.py

@@ -10,6 +10,7 @@ from contextlib import contextmanager
 import torch
 from torch.utils.data import Dataset
 import torchvision.transforms as T
+from PIL import Image
 
 # assumes yolov5 on sys.path
 from lib.yolov5.utils.general import xyxy2xywh
@@ -91,7 +92,7 @@ from torchvision.io import read_image
 import re
 
 class YOLOFrameDataset(Dataset):
-    def __init__(self, img_dir, img_size=896, batch_size=32):
+    def __init__(self, img_dir, img_size=896, batch_size=32, stride=64, pad=0.5):
         self.img_dir = img_dir
         self.img_size = img_size
 
@@ -99,11 +100,41 @@ class YOLOFrameDataset(Dataset):
         self.files = list(filter(lambda f: f[-4:] == ".jpg", self.files))
         self.files.sort(key=lambda f: int(re.sub('\D', '', f)))
 
+        temp_img = read_image(os.path.join(self.img_dir, self.files[0]))
+        size = temp_img.shape
+
+        self.ydim = size[1]
+        self.xdim = size[2]
+
         n = len(self.files)
 
+        aspect_ratio = self.ydim / self.xdim
+        if aspect_ratio < 1:
+            shape = [aspect_ratio, 1]
+        elif aspect_ratio > 1:
+            shape = [1, 1 / aspect_ratio]
+        self.original_shape = (self.ydim, self.xdim)
+        self.shape = np.ceil(np.array(shape) * img_size / stride + pad).astype(int) * stride
+
         self.batches = []
         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
 
     def __len__(self):
         return len(self.batches)
@@ -112,21 +143,28 @@ class YOLOFrameDataset(Dataset):
         for batch_idx in self.batches:
 
             batch = []
+            labels = None
             shapes = []
             for i in range(batch_idx[0], batch_idx[1]):
                 img_name = self.files[i]
                 img_path = os.path.join(self.img_dir, img_name)
-                img = read_image(img_path)
-
-                shapes.append([[img.shape[1], img.shape[2]], None])
-                ratio = self.img_size / max(img.shape[1], img.shape[2])
-                transform = T.Resize([int(ratio*img.shape[1]), int(ratio*img.shape[2])])
-                img = transform(img)
+                image = Image.open(img_path)
+                image = np.asarray(image)
+                img, (h0, w0), (h, w) = self.load_image(image, img_size=self.img_size)
+
+                # Letterbox
+                img, ratio, pad = letterbox(img, self.shape, auto=False, scaleup=False)
+                shape = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling
+                img = img.transpose(2, 0, 1) # to -> C x H x W
+                img = np.ascontiguousarray(img)
+                img = torch.from_numpy(img)
+
+                shapes.append(shape)
                 batch.append(img)
             
             image = torch.stack(batch)
 
-            yield (image, None, shapes)
+            yield (image, labels, shapes)
 
 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,6 +125,7 @@ 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/inferEval.py

@@ -27,7 +27,7 @@ def main(args):
         'min_hits': int(args.min_hits)
     }
 
-    infer(infer_args, config=config)
+    infer(infer_args, config=config, verbose=False)
 
     evaluate("../frames/result_testing", "../frames/MOT", "../frames/metadata", "tracker", True)
 

scripts/infer_frames.py

@@ -39,7 +39,7 @@ def main(args, config={}, verbose=True):
 
     dirname = args.frames
     
-    locations = ["kenai-val"]
+    locations = ["test"]
     for loc in locations:
 
         in_loc_dir = os.path.join(dirname, loc)
@@ -59,9 +59,10 @@ def main(args, config={}, verbose=True):
             for seq in seq_list:
                 pbar.update(1)
                 pbar.set_description("Processing " + seq)
-                print(" ")
-                print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
-                print(" ")
+                if verbose:
+                    print(" ")
+                    print("(" + str(idx) + "/" + str(len(seq_list)) + ") " + seq)
+                    print(" ")
                 idx += 1
                 in_seq_dir = os.path.join(in_loc_dir, seq)
                 infer_seq(in_seq_dir, out_dir, config, seq, model, device, metadata_path, verbose)