Much improved noise augmentation. Simple but effective it seems.

Author: DaWelter

run.sh

@@ -1,31 +1,9 @@
 #!/bin/bash
 
-python -O scripts/train_poseestimator.py --lr 1.e-3 --epochs 1500 --ds "repro_300_wlp+lapa_megaface_lp+wflw_lp" \
+python -O scripts/train_poseestimator.py --lr 1.e-3 --epochs 1500 --ds "repro_300_wlp+lapa_megaface_lp+wflw_lp+replicantface" \
     --with-swa \
     --with-nll-loss \
     --rampup-nll-losses \
     --backbone resnet18 \
-    --outdir model_files/current/baseline4
-
-# python -O scripts/train_poseestimator.py --lr 1.e-3 --epochs 1500 --ds "repro_300_wlp+lapa_megaface_lp+wflw_lp+replicantface" \
-#     --with-swa \
-#     --with-nll-loss \
-#     --rampup-nll-losses \
-#     --backbone resnet18 \
-#     --outdir model_files/current/testing4
-
-# python -O scripts/train_poseestimator.py --lr 1.e-3 --epochs 1500 --ds "repro_300_wlp+lapa_megaface_lp+wflw_lp" \
-#     --with-swa \
-#     --with-nll-loss \
-#     --rampup-nll-losses \
-#     --backbone resnet18 \
-#     --no-pointhead \
-#     --outdir model_files/current/no-pointhead-wo-replicantface4
-
-# python -O scripts/train_poseestimator.py --lr 1.e-3 --epochs 1500 --ds "repro_300_wlp+lapa_megaface_lp+wflw_lp+replicantface" \
-#     --with-swa \
-#     --with-nll-loss \
-#     --rampup-nll-losses \
-#     --backbone resnet18 \
-#     --no-pointhead \
-#     --outdir model_files/current/no-pointhead4
+    --no-pointhead \
+    --outdir model_files/current/more-noise4

trackertraincode/datatransformation/__init__.py

@@ -1,48 +1,3 @@
-# from trackertraincode.datatransformation.tensors import croprescale_image_torch, affine_transform_image_torch
-# from trackertraincode.datatransformation.tensors.image_geometric_cv2 import affine_transform_image_cv2, croprescale_image_cv2
-# from trackertraincode.datatransformation.tensors.affinetrafo import (
-#     position_normalization, position_unnormalization,
-#     apply_affine2d)
-
-# from trackertraincode.datatransformation.batch.intensity import (
-#     KorniaImageDistortions,
-#     RandomBoxBlur,
-#     RandomPlasmaBrightness,
-#     RandomPlasmaContrast,
-#     RandomPlasmaShadow,
-#     RandomGaussianBlur,
-#     RandomSolarize,
-#     RandomInvert,
-#     RandomPosterize,
-#     RandomGamma,
-#     RandomEqualize,
-#     RandomGaussianNoiseWithClipping,
-#     RandomContrast,
-#     RandomBrightness,
-#     RandomGaussianNoise
-# )
-
-# from trackertraincode.datatransformation.batch.normalization import (
-#     normalize_batch,
-#     unnormalize_batch,
-#     offset_points_by_half_pixel,
-#     whiten_batch
-# )
-
-# from trackertraincode.datatransformation.batch.representation import (
-#     to_tensor,
-# )
-
-# from trackertraincode.datatransformation.batch.geometric import (
-#     RandomFocusRoi,
-#     FocusRoi,
-#     RoiFocusRandomizationParameters,
-#     horizontal_flip_and_rot_90
-# )
-
-# from trackertraincode.datatransformation.tensors.normalization import unwhiten_image, whiten_image
-# from trackertraincode.datatransformation.tensors.representation import ensure_image_nchw, ensure_image_nhwc
-
 from . import tensors
 from . import batch
 

trackertraincode/datatransformation/batch/__init__.py

@@ -16,6 +16,7 @@
     RandomContrast,
     RandomBrightness,
     RandomGaussianNoise,
+    OnlyClip,
 )
 
 from .normalization import (

trackertraincode/datatransformation/batch/intensity.py

@@ -23,11 +23,11 @@
     RandomGaussianNoise,
     RandomContrast,
     RandomBrightness,
+    IntensityAugmentationBase2D,
 )
 
 
 class KorniaImageDistortions(object):
-    @wraps(AugmentationSequential.__init__)
     def __init__(self, *args, **kwargs):
         self.augs = AugmentationSequential(*args, **kwargs)
 
@@ -51,3 +51,14 @@ def apply_transform(
         output = super().apply_transform(input, params, flags, transform)
         output = output.clip(0.0, 1.0)
         return output
+
+
+class OnlyClip(IntensityAugmentationBase2D):
+    def apply_transform(
+        self,
+        input: Tensor,
+        params: Dict[str, Tensor],
+        flags: Dict[str, Any],
+        transform: Tensor | None = None,
+    ) -> Tensor:
+        return input.clip(0.0, 1.0)

trackertraincode/pipelines.py

@@ -519,8 +519,11 @@ def make_pose_estimation_loaders(
                 random_apply=4,
             ),
             dtr.batch.KorniaImageDistortions(
-                dtr.batch.RandomGaussianNoise(std=4.0 / 255.0, p=0.5),
-                dtr.batch.RandomGaussianNoise(std=16.0 / 255.0, p=0.1),
+                dtr.batch.RandomGaussianNoise(std=4.0 / 255.0, p=0.25),
+                dtr.batch.RandomGaussianNoise(std=16.0 / 255.0, p=0.25**2),
+                dtr.batch.RandomGaussianNoise(std=32.0 / 255.0, p=0.25**3),
+                dtr.batch.RandomGaussianNoise(std=64.0 / 255.0, p=0.25**4),
+                dtr.batch.OnlyClip(p=1.0),
             ),
         ]
     else: