feat: Added random noise augmentation to object detection

doctr/transforms/modules/pytorch.py

@@ -11,16 +11,16 @@
 from torchvision.transforms import functional as F
 from torchvision.transforms import transforms as T
 
-__all__ = ['Resize']
+__all__ = ['Resize', 'RandomGaussianNoise']
 
 
 class Resize(T.Resize):
     def __init__(
-        self,
-        size: Tuple[int, int],
-        interpolation=F.InterpolationMode.BILINEAR,
-        preserve_aspect_ratio: bool = False,
-        symmetric_pad: bool = False,
+            self,
+            size: Tuple[int, int],
+            interpolation=F.InterpolationMode.BILINEAR,
+            preserve_aspect_ratio: bool = False,
+            symmetric_pad: bool = False,
     ) -> None:
         super().__init__(size, interpolation)
         self.preserve_aspect_ratio = preserve_aspect_ratio
@@ -53,3 +53,15 @@ def __repr__(self) -> str:
         if self.preserve_aspect_ratio:
             _repr += f", preserve_aspect_ratio={self.preserve_aspect_ratio}, symmetric_pad={self.symmetric_pad}"
         return f"{self.__class__.__name__}({_repr})"
+
+
+class RandomGaussianNoise():
+    def __init__(self, mean=0.5, std=1.5):
+        self.std = std
+        self.mean = mean
+
+    def __call__(self, tensor):
+        return tensor + torch.randn(tensor.size()) * self.std + self.mean
+
+    def __repr__(self) -> str:
+        return self.__class__.__name__ + f"mean = {self.mean}, std = {self.std}"

references/obj_detection/train_pytorch.py

@@ -19,6 +19,7 @@
 from fastprogress.fastprogress import master_bar, progress_bar
 from torch.optim.lr_scheduler import MultiplicativeLR, StepLR
 from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
+from torchvision.transforms import Compose
 
 from doctr import transforms as T
 from doctr.datasets import DocArtefacts
@@ -104,33 +105,28 @@ def record_lr(
 
 def convert_to_abs_coords(targets, img_shape):
     height, width = img_shape[-2:]
-    for idx, t in enumerate(targets):
-        targets[idx]['boxes'][:, 0::2] = (t['boxes'][:, 0::2] * width).round()
-        targets[idx]['boxes'][:, 1::2] = (t['boxes'][:, 1::2] * height).round()
-
+    for idx in range(len(targets)):
+        targets[idx]['boxes'][:, 0::2] = (targets[idx]['boxes'][:, 0::2] * width).round()
+        targets[idx]['boxes'][:, 1::2] = (targets[idx]['boxes'][:, 1::2] * height).round()
     targets = [{
         "boxes": torch.from_numpy(t['boxes']).to(dtype=torch.float32),
         "labels": torch.tensor(t['labels']).to(dtype=torch.long)}
         for t in targets
     ]
-
     return targets
 
 
-def fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp=False):
+def fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp):
     if amp:
         scaler = torch.cuda.amp.GradScaler()
-
     model.train()
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
     for images, targets in progress_bar(train_iter, parent=mb):
-
         targets = convert_to_abs_coords(targets, images.shape)
         if torch.cuda.is_available():
             images = images.cuda()
             targets = [{k: v.cuda() for k, v in t.items()} for t in targets]
-
         optimizer.zero_grad()
         if amp:
             with torch.cuda.amp.autocast():
@@ -145,7 +141,6 @@ def fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp=False):
             loss = sum(v for v in loss_dict.values())
             loss.backward()
             optimizer.step()
-
         mb.child.comment = f'Training loss: {loss.item()}'
     scheduler.step()
 
@@ -156,8 +151,8 @@ def evaluate(model, val_loader, metric, amp=False):
     metric.reset()
     val_iter = iter(val_loader)
     for images, targets in val_iter:
-
         images, targets = next(val_iter)
+        # batch_transforms
         targets = convert_to_abs_coords(targets, images.shape)
         if torch.cuda.is_available():
             images = images.cuda()
@@ -238,11 +233,12 @@ def main(args):
         return
 
     st = time.time()
-    # Load both train and val data generators
+
     train_set = DocArtefacts(
         train=True,
         download=True,
-        sample_transforms=T.Resize((args.input_size, args.input_size)),
+        sample_transforms=Compose([T.Resize((args.input_size, args.input_size)),
+                                   T.RandomGaussianNoise(0.5, 1.5)])
     )
 
     train_loader = DataLoader(
@@ -305,7 +301,6 @@ def main(args):
 
     mb = master_bar(range(args.epochs))
     max_score = 0.
-
     for epoch in mb:
         fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp=args.amp)
         # Validation loop at the end of each epoch
@@ -340,7 +335,7 @@ def parse_args():
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
     parser.add_argument('arch', type=str, help='text-detection model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
-    parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
+    parser.add_argument('--epochs', type=int, default=20, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=2, help='batch size for training')
     parser.add_argument('--device', default=None, type=int, help='device')
     parser.add_argument('--input_size', type=int, default=1024, help='model input size, H = W')

references/obj_detection/utils.py

@@ -3,6 +3,7 @@
 # This program is licensed under the Apache License version 2.
 # See LICENSE or go to <https://www.apache.org/licenses/LICENSE-2.0.txt> for full license details.
 
+
 from typing import Dict, List
 
 import cv2