Add code to check inputs to all loss functions are valid

odak/learn/perception/blur_loss.py

@@ -1,6 +1,7 @@
 import torch
 
 from .radially_varying_blur import RadiallyVaryingBlur
+from .util import check_loss_inputs
 
 
 class BlurLoss():
@@ -68,6 +69,7 @@ def __call__(self, image, target, gaze=[0.5, 0.5]):
         loss                : torch.tensor
                                 The computed loss.
         """
+        check_loss_inputs("BlurLoss", image, target)
         blurred_target = self.blur_image(target, gaze)
         if self.blur_source:
             blurred_image = self.blur_image(image, gaze)

odak/learn/perception/metamer_mse_loss.py

@@ -1,9 +1,9 @@
 import torch
-import math
 
 from .metameric_loss import MetamericLoss
 from .color_conversion import ycrcb_2_rgb, rgb_2_ycrcb
 from .spatial_steerable_pyramid import pad_image_for_pyramid
+from .util import check_loss_inputs
 
 
 class MetamerMSELoss():
@@ -128,6 +128,7 @@ def __call__(self, image, target, gaze=[0.5, 0.5]):
         loss                : torch.tensor
                                 The computed loss.
         """
+        check_loss_inputs("MetamerMSELoss", image, target)
         # Pad image and target if necessary
         image = pad_image_for_pyramid(image, self.metameric_loss.n_pyramid_levels)
         target = pad_image_for_pyramid(target, self.metameric_loss.n_pyramid_levels)

odak/learn/perception/metameric_loss.py

@@ -1,11 +1,10 @@
-from os import stat
 import torch
-import math
 
-from .color_conversion import ycrcb_2_rgb, rgb_2_ycrcb
+from .color_conversion import rgb_2_ycrcb
 from .spatial_steerable_pyramid import SpatialSteerablePyramid, pad_image_for_pyramid
 from .radially_varying_blur import RadiallyVaryingBlur
 from .foveation import make_radial_map
+from .util import check_loss_inputs
 
 
 class MetamericLoss():
@@ -208,13 +207,11 @@ def __call__(self, image, target, gaze=[0.5, 0.5], image_colorspace="RGB", visua
         loss                : torch.tensor
                                 The computed loss.
         """
-        # TODO if input is RGB, convert to YCrCb.
-        if image.size(1) != target.size(1):
-            raise Exception(
-                "MetamericLoss ERROR: Input and target must have same number of channels.")
+        check_loss_inputs("MetamericLoss", image, target)
         # Pad image and target if necessary
         image = pad_image_for_pyramid(image, self.n_pyramid_levels)
         target = pad_image_for_pyramid(target, self.n_pyramid_levels)
+        # If input is RGB, convert to YCrCb.
         if image.size(1) == 3 and image_colorspace == "RGB":
             image = rgb_2_ycrcb(image)
             target = rgb_2_ycrcb(target)

odak/learn/perception/metameric_loss_uniform.py

@@ -1,11 +1,9 @@
-from os import stat
 import torch
-import math
 
-from .color_conversion import ycrcb_2_rgb, rgb_2_ycrcb
+from .color_conversion import rgb_2_ycrcb
 from .spatial_steerable_pyramid import SpatialSteerablePyramid, pad_image_for_pyramid
-from .radially_varying_blur import RadiallyVaryingBlur
-from .foveation import make_radial_map
+from .util import check_loss_inputs
+
 
 def uniform_blur(image, pooling_size):
     original_size = image.size()[-2:]
@@ -127,13 +125,11 @@ def __call__(self, image, target, image_colorspace="RGB", visualise_loss=False):
         loss                : torch.tensor
                                 The computed loss.
         """
-        # TODO if input is RGB, convert to YCrCb.
-        if image.size(1) != target.size(1):
-            raise Exception(
-                "MetamericLoss ERROR: Input and target must have same number of channels.")
+        check_loss_inputs("MetamericLossUniform", image, target)
         # Pad image and target if necessary
         image = pad_image_for_pyramid(image, self.n_pyramid_levels)
         target = pad_image_for_pyramid(target, self.n_pyramid_levels)
+        # If input is RGB, convert to YCrCb.
         if image.size(1) == 3 and image_colorspace == "RGB":
             image = rgb_2_ycrcb(image)
             target = rgb_2_ycrcb(target)

odak/learn/perception/radially_varying_blur.py

@@ -1,8 +1,4 @@
-import numpy as np
-import sys
 import torch
-from math import floor
-import math
 
 from .foveation import make_pooling_size_map_lod
 

odak/learn/perception/spatial_steerable_pyramid.py

@@ -1,8 +1,8 @@
 from .steerable_pyramid_filters import get_steerable_pyramid_filters
 import torch
-import numpy as np
 import math
 
+
 def pad_image_for_pyramid(image, n_pyramid_levels):
     """
     Pads an image to the extent necessary to compute a steerable pyramid of the input image.

odak/learn/perception/steerable_pyramid_filters.py

@@ -1,5 +1,4 @@
 import torch
-from torch.functional import norm
 
 
 def crop_steerable_pyramid_filters(filters, size):

odak/learn/perception/util.py

@@ -0,0 +1,15 @@
+import torch
+
+
+def check_loss_inputs(loss_name, image, target):
+        if image.size() != target.size():
+            raise Exception(
+                loss_name + " ERROR: Input and target must have same dimensions.")
+        if len(image.size()) != 4:
+            raise Exception(
+                loss_name + " ERROR: Input and target must have 4 dimensions in total (NCHW format).")
+        if image.size(1) != 1 and image.size(1) != 3:
+            raise Exception(
+                loss_name + """ ERROR: Inputs should have either 1 or 3 channels 
+                (1 channel for grayscale, 3 for RGB or YCrCb).
+                Ensure inputs have 1 or 3 channels and are in NCHW format.""")