models/create_fasterrcnn_model.py

datasets.py

@@ -11,6 +11,7 @@
     get_train_transform, 
     get_valid_transform,
     get_train_aug,
+    get_train_aug_custom,
     transform_mosaic
 )
 
@@ -23,11 +24,13 @@ def __init__(
         labels_path, 
         img_size, 
         classes, 
+        aug_option,
         transforms=None, 
         use_train_aug=False,
         train=False, 
         no_mosaic=False,
-        square_training=False
+        square_training=False,
+
     ):
         self.transforms = transforms
         self.use_train_aug = use_train_aug
@@ -43,6 +46,7 @@ def __init__(
         self.all_image_paths = []
         self.log_annot_issue_x = True
         self.log_annot_issue_y = True
+        self.aug_option=aug_option
 
         # get all the image paths in sorted order
         for file_type in self.image_file_types:
@@ -316,13 +320,14 @@ def __getitem__(self, idx):
         target["image_id"] = image_id
 
 
-        if self.use_train_aug: # Use train augmentation if argument is passed.
-            train_aug = get_train_aug()
-            sample = train_aug(image=image_resized,
-                                     bboxes=target['boxes'],
-                                     labels=labels)
-            image_resized = sample['image']
-            target['boxes'] = torch.Tensor(sample['bboxes']).to(torch.int64)
+        if self.use_train_aug: # Use train augmentation if argument is passed.    
+                train_aug = get_train_aug_custom(self.aug_option)
+                sample = train_aug(image=image_resized,
+                                        bboxes=target['boxes'],
+                                        labels=labels)
+                image_resized = sample['image']
+                target['boxes'] = torch.Tensor(sample['bboxes']).to(torch.int64)
+
         else:
             sample = self.transforms(image=image_resized,
                                      bboxes=target['boxes'],
@@ -352,6 +357,7 @@ def create_train_dataset(
     train_dir_labels, 
     img_size, 
     classes,
+    aug_option,
     use_train_aug=False,
     no_mosaic=False,
     square_training=False
@@ -361,6 +367,7 @@ def create_train_dataset(
         train_dir_labels,
         img_size, 
         classes, 
+        aug_option,
         get_train_transform(),
         use_train_aug=use_train_aug,
         train=True, 
@@ -373,13 +380,15 @@ def create_valid_dataset(
     valid_dir_labels, 
     img_size, 
     classes,
+    aug_option,
     square_training=False
 ):
     valid_dataset = CustomDataset(
         valid_dir_images, 
         valid_dir_labels, 
         img_size, 
         classes, 
+        aug_option,
         get_valid_transform(),
         train=False, 
         no_mosaic=True,
@@ -411,4 +420,4 @@ def create_valid_loader(
         collate_fn=collate_fn,
         sampler=batch_sampler
     )
-    return valid_loader
+    return valid_loader

models/create_fasterrcnn_model.py

@@ -16,6 +16,13 @@ def return_fasterrcnn_mobilenetv3_large_fpn(
     )
     return model
 
+def return_fasterrcnn_efficientnet_b4(
+    num_classes, pretrained=True, coco_model=False
+):
+    model = fasterrcnn_efficientnet_b4.create_model(
+        num_classes, pretrained=pretrained, coco_model=coco_model
+    )
+    return model
 def return_fasterrcnn_mobilenetv3_large_320_fpn(
     num_classes, pretrained=True, coco_model=False
 ):    
@@ -200,9 +207,18 @@ def return_fasterrcnn_regnet_y_400mf(
     )
     return model
 
+def return_fasterrcnn_mbv3_large(
+    num_classes, pretrained=True, coco_model=False
+):
+    model = fasterrcnn_mbv3_large.create_model(
+        num_classes, pretrained, coco_model=coco_model
+    )
+    return model
+
 create_model = {
     'fasterrcnn_resnet50_fpn': return_fasterrcnn_resnet50_fpn,
     'fasterrcnn_mobilenetv3_large_fpn': return_fasterrcnn_mobilenetv3_large_fpn,
+    'fasterrcnn_efficientnet_b4':return_fasterrcnn_efficientnet_b4,
     'fasterrcnn_mobilenetv3_large_320_fpn': return_fasterrcnn_mobilenetv3_large_320_fpn,
     'fasterrcnn_resnet18': return_fasterrcnn_resnet18,
     'fasterrcnn_custom_resnet': return_fasterrcnn_custom_resnet,
@@ -225,5 +241,6 @@ def return_fasterrcnn_regnet_y_400mf(
     'fasterrcnn_vitdet': return_fasterrcnn_vitdet,
     'fasterrcnn_vitdet_tiny': return_fasterrcnn_vitdet_tiny,
     'fasterrcnn_mobilevit_xxs': return_fasterrcnn_mobilevit_xxs,
-    'fasterrcnn_regnet_y_400mf': return_fasterrcnn_regnet_y_400mf
-}
+    'fasterrcnn_regnet_y_400mf': return_fasterrcnn_regnet_y_400mf,
+    'fasterrcnn_mbv3_large':return_fasterrcnn_mbv3_large,
+}

requirements.txt

@@ -3,8 +3,8 @@ albumentations>=1.1.0
 ipython
 jupyter
 matplotlib
-opencv-python>=4.1.1.26
-opencv-python-headless>=4.1.1.26
+opencv-python~=4.1
+opencv-python-headless~=4.1
 Pillow
 PyYAML
 scikit-image
@@ -23,7 +23,8 @@ tensorboard
 
 # Model summary----------------------------
 torchinfo
-
+# Transformer based models.
+vision_transformers
 # Extras-----------------------------------
 pycocotools>=2.0.2
 setuptools==59.5.0
@@ -32,4 +33,4 @@ torchmetrics # Evaluation
 # Export-----------------------------------
 # onnxruntime
 # onnx
-# onnxruntime-gpu
+# onnxruntime-gpu

setup.cfg

@@ -0,0 +1,20 @@
+[metadata]
+name = fasterrcnn_pytorch_training_pipeline
+version = 0.1.0
+author = Sovit Ranjan Rath
+author_email = [email protected]
+description = A Simple Pipeline to Train PyTorch FasterRCNN Model
+long_description = file: README.md
+url = https://github.com/falibabaei/fasterrcnn_pytorch_training_pipeline.git
+license = MIT
+
+classifier =
+    Intended Audience :: Information Technology
+    Intended Audience :: System Administrators
+    License :: OSI Approved :: MIT License
+    Operating System :: POSIX :: Linux
+    Programming Language :: Python
+    Programming Language :: Python :: 3
+    Programming Language :: Python :: 3.6
+
+

setup.py

@@ -0,0 +1,30 @@
+from setuptools import setup, find_packages
+
+with open("README.md", "r") as fh:
+    long_description = fh.read()
+
+with open("requirements.txt", "r") as fh:
+    requirements = [line.strip() for line in fh]
+
+setup(
+    name='fasterrcnn_pytorch_training_pipeline',
+    version='0.1.0',
+    author='Sovit Ranjan Rath',
+    author_email='[email protected]',
+    description='A Simple Pipeline to Train PyTorch FasterRCNN Model',
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url='https://github.com/falibabaei/fasterrcnn_pytorch_training_pipeline.git',
+    license='MIT',
+    classifiers=[
+        'Intended Audience :: Information Technology',
+        'Intended Audience :: System Administrators',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: POSIX :: Linux',
+        'Programming Language :: Python',
+        'Programming Language :: Python :: 3',
+        'Programming Language :: Python :: 3.6',
+    ],
+    packages=find_packages(),
+    install_requires=requirements
+)

torch_utils/engine.py

@@ -158,6 +158,9 @@ def evaluate(
             val_saved_image = save_validation_results(
                 images, outputs, counter, out_dir, classes, colors
             )
+        else:
+            val_saved_image=None    
+
 
     # gather the stats from all processes
     metric_logger.synchronize_between_processes()
@@ -168,4 +171,4 @@ def evaluate(
     coco_evaluator.accumulate()
     stats = coco_evaluator.summarize()
     torch.set_num_threads(n_threads)
-    return stats, val_saved_image
+    return stats, val_saved_image

utils/logging.py

@@ -181,44 +181,44 @@ def wandb_log(
     )
 
     bg = np.full((image_size * 2, image_size * 2, 3), 114, dtype=np.float32)
-
-    if len(val_pred_image) == 1:
-        log_image = overlay_on_canvas(bg, val_pred_image[0])
-        wandb.log({'predictions': [wandb.Image(log_image)]})
-
-    if len(val_pred_image) == 2:
-        log_image = cv2.hconcat(
-            [
-                overlay_on_canvas(bg, val_pred_image[0]), 
-                overlay_on_canvas(bg, val_pred_image[1])
-            ]
-        )
-        wandb.log({'predictions': [wandb.Image(log_image)]})
-
-    if len(val_pred_image) > 2 and len(val_pred_image) <= 8:
-        log_image = overlay_on_canvas(bg, val_pred_image[0])
-        for i in range(len(val_pred_image)-1):
-            log_image = cv2.hconcat([
-                log_image, 
-                overlay_on_canvas(bg, val_pred_image[i+1])
-            ])
-        wandb.log({'predictions': [wandb.Image(log_image)]})
+    if val_pred_image is not None:
+        if   len(val_pred_image) == 1:
+            log_image = overlay_on_canvas(bg, val_pred_image[0])
+            wandb.log({'predictions': [wandb.Image(log_image)]})
+
+        elif len(val_pred_image) == 2:
+            log_image = cv2.hconcat(
+                [
+                    overlay_on_canvas(bg, val_pred_image[0]), 
+                    overlay_on_canvas(bg, val_pred_image[1])
+                ]
+            )
+            wandb.log({'predictions': [wandb.Image(log_image)]})
 
-    if len(val_pred_image) > 8:
-        log_image = overlay_on_canvas(bg, val_pred_image[0])
-        for i in range(len(val_pred_image)-1):
-            if i == 7:
-                break
-            log_image = cv2.hconcat([
-                log_image, 
-                overlay_on_canvas(bg, val_pred_image[i-1])
-            ])
-        wandb.log({'predictions': [wandb.Image(log_image)]})
+        elif len(val_pred_image) > 2 and len(val_pred_image) <= 8:
+            log_image = overlay_on_canvas(bg, val_pred_image[0])
+            for i in range(len(val_pred_image)-1):
+                log_image = cv2.hconcat([
+                    log_image, 
+                    overlay_on_canvas(bg, val_pred_image[i+1])
+                ])
+            wandb.log({'predictions': [wandb.Image(log_image)]})
+
+        elif len(val_pred_image) > 8:
+            log_image = overlay_on_canvas(bg, val_pred_image[0])
+            for i in range(len(val_pred_image)-1):
+                if i == 7:
+                    break
+                log_image = cv2.hconcat([
+                    log_image, 
+                    overlay_on_canvas(bg, val_pred_image[i-1])
+                ])
+            wandb.log({'predictions': [wandb.Image(log_image)]})
 
 def wandb_save_model(model_dir):
     """
     Uploads the models to Weights&Biases.
 
     :param model_dir: Local disk path where models are saved.
     """
-    wandb.save(os.path.join(model_dir, 'best_model.pth'))
+    wandb.save(os.path.join(model_dir, 'best_model.pth'))

utils/transforms.py

@@ -31,7 +31,45 @@ def get_train_aug():
         'format': 'pascal_voc',
         'label_fields': ['labels'],
     })
+def get_train_aug_custom(augmentation_choices):
+    transform_list = []
+    for key in augmentation_choices.keys():
+        if key == 'blur':
+            transform_list.append(A.Blur(p=augmentation_choices[key]['p'],blur_limit=augmentation_choices[key]['blur_limit']))
+        elif key == 'motion_blur':
+            transform_list.append(A.MotionBlur(p=augmentation_choices[key]['p'],blur_limit=augmentation_choices[key]['blur_limit']))
+        elif key == 'median_blur':
+            transform_list.append(A.MedianBlur(p=augmentation_choices[key]['p'],blur_limit=augmentation_choices[key]['blur_limit']))
+        elif key == 'to_gray':
+            transform_list.append(A.ToGray(p=augmentation_choices[key]['p']))
+        elif key == 'random_brightness_contrast':
+            transform_list.append(A.RandomBrightnessContrast(p=augmentation_choices[key]['p']))
+        elif key == 'color_jitter':
+            transform_list.append(A.ColorJitter(p=augmentation_choices[key]['p']))
+        elif key == 'random_gamma':
+            transform_list.append(A.RandomGamma(p=augmentation_choices[key]['p']))
+        elif key == 'horizontal_flip':
+            transform_list.append(A.HorizontalFlip(p=augmentation_choices[key]['p']))
+        elif key == 'vertical_flip':
+            transform_list.append(A.VerticalFlip(p=augmentation_choices[key]['p']))
+        elif key == 'rotate':
+            transform_list.append(A.Rotate(limit=augmentation_choices[key]['p']))
+        elif key == 'shift_scale_rotate':
+            transform_list.append(A.ShiftScaleRotate(shift_limit=augmentation_choices[key]['shift_limit'], 
+                                                     scale_limit=augmentation_choices[key]['scale_limit'], 
+                                                     rotate_limit=augmentation_choices[key]['rotate_limit']))
+        elif key == 'Cutout':
+            transform_list.append(A.ShiftScaleRotate(num_holes=augmentation_choices[key]['num_holes'], 
+                                                     max_h_size=augmentation_choices[key]['max_h_size'], 
+                                                     fill_value=augmentation_choices[key]['fill_value'],
+                                                     p=augmentation_choices[key]['p']))    
+        elif key == 'rotate':
+            transform_list.append(A.ChannelShuffle(p=augmentation_choices[key]['p']))    
+        transform_list.append(ToTensorV2(p=1.0))
 
+        bbox_params = {'format': 'pascal_voc', 'label_fields': ['labels']}
+
+        return A.Compose(transform_list, bbox_params=bbox_params)
 def get_train_transform():
     return A.Compose([
         ToTensorV2(p=1.0),
@@ -86,4 +124,4 @@ def infer_transforms(image):
         transforms.ToPILImage(),
         transforms.ToTensor(),
     ])
-    return transform(image)
+    return transform(image)