Add wrapped MSE loss and unit test for loss.py

odak/learn/models/models.py

@@ -140,8 +140,8 @@ def __init__(
                                       kernel_size = kernel_size,
                                       bias = bias,
                                       activation = activation
-                                     )
-        factor = 2 if bilinear else 1        
+                                     )      
+
         self.downsampling_layers = torch.nn.ModuleList()
         self.upsampling_layers = torch.nn.ModuleList()
         for i in range(depth): # downsampling layers

odak/learn/tools/loss.py

@@ -121,18 +121,28 @@ def histogram_loss(frame, ground_truth, bins = 32, limits = [0., 1.]):
                        Loss from evaluation.
     """
     if len(frame.shape) == 2:
+        frame = frame.unsqueeze(0).unsqueeze(0)
+    elif len(frame.shape) == 3:
         frame = frame.unsqueeze(0)
-    if len(frame.shape) == 3:
-        frame = frame.unsqueeze(0)
+
+    if len(ground_truth.shape) == 2:
+        ground_truth = ground_truth.unsqueeze(0).unsqueeze(0)
+    elif len(ground_truth.shape) == 3:
+        ground_truth = ground_truth.unsqueeze(0)
+
     histogram_frame = torch.zeros(frame.shape[1], bins).to(frame.device)
-    histogram_ground_truth = torch.zeros(frame.shape[1], bins).to(frame.device)
+    histogram_ground_truth = torch.zeros(ground_truth.shape[1], bins).to(frame.device)
+
     l2 = torch.nn.MSELoss()
+
     for i in range(frame.shape[1]):
-        histogram_frame[i] = torch.histc(frame[:, i].flatten(), bins = bins, min = limits[0], max = limits[1])
-        histogram_ground_truth[i] = torch.histc(frame[:, i].flatten(), bins = bins, min = limits[0], max = limits[1])
+        histogram_frame[i] = torch.histc(frame[:, i].flatten(), bins=bins, min=limits[0], max=limits[1])
+        histogram_ground_truth[i] = torch.histc(ground_truth[:, i].flatten(), bins=bins, min=limits[0], max=limits[1])
+
     loss = l2(histogram_frame, histogram_ground_truth)
-    return loss
 
+    return loss
+
 
 def weber_contrast(image, roi_high, roi_low):
     """
@@ -141,7 +151,7 @@ def weber_contrast(image, roi_high, roi_low):
     Parameters
     ----------
     image         : torch.tensor
-                    Image to be tested [1 x 3 x m x n] or [3 x m x n] or [m x n].
+                    Image to be tested [1 x 3 x m x n] or [3 x m x n] or [1 x m x n] or [m x n].
     roi_high      : torch.tensor
                     Corner locations of the roi for high intensity area [m_start, m_end, n_start, n_end].
     roi_low       : torch.tensor
@@ -192,3 +202,33 @@ def michelson_contrast(image, roi_high, roi_low):
     low = torch.mean(region_low, dim = (2, 3))
     result = (high - low) / (high + low)
     return result.squeeze(0)
+
+
+def wrapped_mean_squared_error(image, ground_truth, reduction = 'mean'):
+    """
+    A function to calculate the wrapped mean squared error between predicted and target angles.
+
+    Parameters
+    ----------
+    image         : torch.tensor
+                    Image to be tested [1 x 3 x m x n]  or [3 x m x n] or [1 x m x n] or [m x n].
+    ground_truth  : torch.tensor
+                    Ground truth to be tested [1 x 3 x m x n]  or [3 x m x n] or [1 x m x n] or [m x n].
+    reduction     : str
+                    Specifies the reduction to apply to the output: 'mean' (default) or 'sum'.
+
+    Returns
+    -------
+    wmse        : torch.tensor
+                  The calculated wrapped mean squared error. 
+    """
+    sin_diff = torch.sin(image) - torch.sin(ground_truth)
+    cos_diff = torch.cos(image) - torch.cos(ground_truth)
+    loss = (sin_diff**2 + cos_diff**2)
+
+    if reduction == 'mean':
+        return loss.mean()
+    elif reduction == 'sum':
+        return loss.sum()
+    else:
+        raise ValueError("Invalid reduction type. Choose 'mean' or 'sum'.")

test/test_learn_components.py

@@ -5,7 +5,7 @@
 def test():
     # test residual block
     x = torch.randn(1, 2, 32, 32)
-    residual_inference= components.residual_layer()
+    residual_inference = components.residual_layer()
     y = residual_inference(x)
     # test convolution layer
     convolution_inference = components.convolution_layer()

test/test_tools_losses.py

@@ -0,0 +1,26 @@
+import torch
+import sys
+import odak.learn.tools.loss as loss
+
+
+def test():
+    # test residual block
+    image = [torch.randn(1, 3, 32, 32), torch.randn(1, 32, 32), torch.randn(32, 32), torch.randn(3, 32, 32)]
+    ground_truth = [torch.randn(1, 3, 32, 32), torch.randn(1, 32, 32), torch.randn(32, 32), torch.randn(3, 32, 32)]
+    for idx, (img, pred) in enumerate(zip(image, ground_truth)):
+        print(f'Running test {idx}, input shape: {img.size()}...')
+        y = loss.psnr(img, pred, peak_value = 2.0)
+        y = loss.multi_scale_total_variation_loss(img, levels = 4)
+        y = loss.total_variation_loss(img)
+        y = loss.histogram_loss(img, pred, bins = 16, limits = [0., 1.])
+        roi_high = [0, 16, 0, 16]
+        roi_low = [16, 32, 16, 32]
+        y = loss.weber_contrast(img, roi_high, roi_low)
+        y = loss.michelson_contrast(img, roi_high, roi_low)
+        y = loss.wrapped_mean_squared_error(img, pred, reduction='sum')
+    value = torch.tensor(1.0, dtype=torch.float)
+    y = loss.radial_basis_function(value = value, epsilon = 0.5)
+    assert True == True
+
+if __name__ == '__main__':
+    sys.exit(test())