🔨 [Update] dataloader, return data augment info

yolo/tools/data_augmentation.py

@@ -10,7 +10,7 @@ class AugmentationComposer:
     def __init__(self, transforms, image_size: int = [640, 640]):
         self.transforms = transforms
         # TODO: handle List of image_size [640, 640]
-        self.image_size = image_size[0]
+        self.image_size = image_size
         self.pad_resize = PadAndResize(self.image_size)
 
         for transform in self.transforms:
@@ -29,27 +29,32 @@ class AugmentationComposer:
 
 
 class PadAndResize:
-    def __init__(self, image_size):
+    def __init__(self, image_size, background_color=(128, 128, 128)):
         """Initialize the object with the target image size."""
-        self.image_size = image_size
+        self.target_width, self.target_height = image_size
+        self.background_color = background_color
 
-    def __call__(self, image, boxes):
-        original_size = max(image.size)
-        scale = self.image_size / original_size
-        square_img = Image.new("RGB", (original_size, original_size), (128, 128, 128))
-        left = (original_size - image.width) // 2
-        top = (original_size - image.height) // 2
-        square_img.paste(image, (left, top))
+    def __call__(self, image: Image, boxes):
+        img_width, img_height = image.size
+        scale = min(self.target_width / img_width, self.target_height / img_height)
+        new_width, new_height = int(img_width * scale), int(img_height * scale)
+
+        resized_image = image.resize((new_width, new_height), Image.LANCZOS)
 
-        resized_img = square_img.resize((self.image_size, self.image_size))
+        pad_left = (self.target_width - new_width) // 2
+        pad_top = (self.target_height - new_height) // 2
+        padded_image = Image.new("RGB", (self.target_width, self.target_height), self.background_color)
+        padded_image.paste(resized_image, (pad_left, pad_top))
 
-        boxes[:, 1] = (boxes[:, 1] * image.width + left) / self.image_size * scale
-        boxes[:, 2] = (boxes[:, 2] * image.height + top) / self.image_size * scale
-        boxes[:, 3] = (boxes[:, 3] * image.width + left) / self.image_size * scale
-        boxes[:, 4] = (boxes[:, 4] * image.height + top) / self.image_size * scale
+        boxes[:, 1] *= scale  # xmin
+        boxes[:, 2] *= scale  # ymin
+        boxes[:, 3] *= scale  # xmax
+        boxes[:, 4] *= scale  # ymax
+        boxes[:, [1, 3]] += pad_left
+        boxes[:, [2, 4]] += pad_top
 
-        rev_tensor = torch.tensor([scale, left, top, left, top])
-        return resized_img, boxes, rev_tensor
+        transform_info = torch.tensor([scale, pad_left, pad_top, pad_left, pad_top])
+        return padded_image, boxes, transform_info
 
 
 class HorizontalFlip:
@@ -94,7 +99,7 @@ class Mosaic:
 
         assert self.parent is not None, "Parent is not set. Mosaic cannot retrieve image size."
 
-        img_sz = self.parent.image_size  # Assuming `image_size` is defined in parent
+        img_sz = self.parent.image_size[0]  # Assuming `image_size` is defined in parent
         more_data = self.parent.get_more_data(3)  # get 3 more images randomly
 
         data = [(image, boxes)] + more_data

yolo/tools/data_loader.py

@@ -141,16 +141,16 @@ class YoloDataset(Dataset):
     def get_data(self, idx):
         img_path, bboxes = self.data[idx]
         img = Image.open(img_path).convert("RGB")
-        return img, bboxes
+        return img, bboxes, img_path
 
     def get_more_data(self, num: int = 1):
         indices = torch.randint(0, len(self), (num,))
-        return [self.get_data(idx) for idx in indices]
+        return [self.get_data(idx)[:2] for idx in indices]
 
     def __getitem__(self, idx) -> Union[Image.Image, torch.Tensor]:
-        img, bboxes = self.get_data(idx)
-        img, bboxes, _ = self.transform(img, bboxes)
-        return img, bboxes
+        img, bboxes, img_path = self.get_data(idx)
+        img, bboxes, rev_tensor = self.transform(img, bboxes)
+        return img, bboxes, rev_tensor, img_path
 
     def __len__(self) -> int:
         return len(self.data)
@@ -195,9 +195,11 @@ class YoloDataLoader(DataLoader):
             batch_targets[idx, :target_size] = batch[idx][1]
         batch_targets[:, :, 1:] *= self.image_size
 
-        batch_images = torch.stack([item[0] for item in batch])
+        batch_images, _, batch_reverse, batch_path = zip(*batch)
+        batch_images = torch.stack(batch_images)
+        batch_reverse = torch.stack(batch_reverse)
 
-        return batch_images, batch_targets
+        return batch_images, batch_targets, batch_reverse, batch_path
 
 
 def create_dataloader(data_cfg: DataConfig, dataset_cfg: DatasetConfig, task: str = "train", use_ddp: bool = False):
@@ -261,12 +263,14 @@ class StreamDataLoader:
         if isinstance(frame, np.ndarray):
             frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
             frame = Image.fromarray(frame)
+        origin_frame = frame
         frame, _, rev_tensor = self.transform(frame, torch.zeros(0, 5))
         frame = frame[None]
+        rev_tensor = rev_tensor[None]
         if not self.is_stream:
-            self.queue.put(frame)
+            self.queue.put((frame, rev_tensor, origin_frame))
         else:
-            self.current_frame = frame
+            self.current_frame = (frame, rev_tensor, origin_frame)
 
     def __iter__(self) -> Generator[Tensor, None, None]:
         return self

yolo/tools/solver.py

@@ -72,7 +72,7 @@ class ModelTrainer:
         self.model.train()
         total_loss = 0
 
-        for images, targets in dataloader:
+        for images, targets, *_ in dataloader:
             loss, loss_each = self.train_one_batch(images, targets)
 
             total_loss += loss
@@ -136,8 +136,9 @@ class ModelTester:
 
             last_time = time.time()
         try:
-            for idx, images in enumerate(dataloader):
+            for idx, (images, rev_tensor, origin_frame) in enumerate(dataloader):
                 images = images.to(self.device)
+                rev_tensor = rev_tensor.to(self.device)
                 with torch.no_grad():
                     predicts = self.model(images)
                     predicts = self.vec2box(predicts["Main"])
@@ -192,8 +193,8 @@ class ModelValidator:
         iou_thresholds = torch.arange(0.5, 1.0, 0.05)
         map_all = []
         self.progress.start_one_epoch(len(dataloader))
-        for images, targets in dataloader:
-            images, targets = images.to(self.device), targets.to(self.device)
+        for images, targets, rev_tensor, img_paths in dataloader:
+            images, targets, rev_tensor = images.to(self.device), targets.to(self.device), rev_tensor.to(self.device)
             with torch.no_grad():
                 predicts = self.model(images)
             predicts = self.vec2box(predicts["Main"])