image_to_image.losses.weighted_combined_loss

  1"""
  2Module to define a Weighted Combine Loss.
  3
  4Functions:
  5- calc_weight_map
  6
  7Classes:
  8- WeightedCombinedLoss
  9
 10By Tobia Ippolito
 11"""
 12# ---------------------------
 13#        > Imports <
 14# ---------------------------
 15import torch
 16import torch.nn as nn
 17import torch.nn.functional as F
 18
 19import kornia
 20
 21
 22
 23# ---------------------------
 24#   > Loss Implementation <
 25# ---------------------------
 26class WeightedCombinedLoss(nn.Module):
 27    """
 28    Computes a weighted combination of multiple loss functions for image-to-image tasks.
 29
 30    Supported losses:
 31        - SILog loss
 32        - Gradient L1 loss
 33        - SSIM loss
 34        - Edge-aware loss
 35        - L1 loss
 36        - Variance loss
 37        - Range loss
 38        - Blur loss
 39
 40    The losses can be weighted individually, and average losses are tracked across steps.
 41    """
 42    def __init__(self, 
 43                 silog_lambda=0.5, 
 44                 weight_silog=0.5, 
 45                 weight_grad=10.0, 
 46                 weight_ssim=5.0,
 47                 weight_edge_aware=10.0,
 48                 weight_l1=1.0,
 49                 weight_var=1.0,
 50                 weight_range=1.0,
 51                 weight_blur=1.0):
 52        """
 53        Initializes the WeightedCombinedLoss with optional weights for each component.
 54
 55        Parameter:
 56        - silog_lambda (float): 
 57            SILog lambda parameter.
 58        - weight_silog (float): 
 59            Weight for SILog loss.
 60        - weight_grad (float): 
 61            Weight for gradient L1 loss.
 62        - weight_ssim (float): 
 63            Weight for SSIM loss.
 64        - weight_edge_aware (float): 
 65            Weight for edge-aware loss.
 66        - weight_l1 (float): 
 67            Weight for L1 loss.
 68        - weight_var (float): 
 69            Weight for variance loss.
 70        - weight_range (float): 
 71            Weight for range loss.
 72        - weight_blur (float): 
 73            Weight for blur loss.
 74        """
 75        super().__init__()
 76        self.silog_lambda = silog_lambda
 77        self.weight_silog = weight_silog
 78        self.weight_grad = weight_grad
 79        self.weight_ssim = weight_ssim
 80        self.weight_edge_aware = weight_edge_aware
 81        self.weight_l1 = weight_l1
 82        self.weight_var = weight_var
 83        self.weight_range = weight_range
 84        self.weight_blur = weight_blur
 85
 86        self.avg_loss_silog = 0
 87        self.avg_loss_grad = 0
 88        self.avg_loss_ssim = 0
 89        self.avg_loss_l1 = 0
 90        self.avg_loss_edge_aware = 0
 91        self.avg_loss_var = 0
 92        self.avg_loss_range = 0
 93        self.avg_loss_blur = 0
 94        self.steps = 0
 95
 96        # Instantiate SSIMLoss module
 97        self.ssim_module = kornia.losses.SSIMLoss(window_size=11, reduction='mean')
 98        # self.ssim_module = kornia.losses.MS_SSIMLoss(reduction='mean')
 99
100
101    def silog_loss(self, pred, target, weight_map):
102        eps = 1e-6
103        pred = torch.clamp(pred, min=eps)
104        target = torch.clamp(target, min=eps)
105        
106        diff_log = torch.log(target) - torch.log(pred)
107        diff_log = diff_log * weight_map
108
109        loss = torch.sqrt(torch.mean(diff_log ** 2) -
110                          self.silog_lambda * torch.mean(diff_log) ** 2)
111        return loss
112
113    def gradient_l1_loss(self, pred, target, weight_map):
114        # Create Channel Dimension
115        if pred.ndim == 3:
116            pred = pred.unsqueeze(1)
117        if target.ndim == 3:
118            target = target.unsqueeze(1)
119        if weight_map.ndim == 3:
120            weight_map = weight_map.unsqueeze(1)
121
122        # Gradient in x-direction (horizontal -> dim=3)
123        pred_grad_x = pred[:, :, :, 1:] - pred[:, :, :, :-1]
124        target_grad_x = target[:, :, :, 1:] - target[:, :, :, :-1]
125
126        # Gradient in y-direction (vertical -> dim=2)
127        pred_grad_y = pred[:, :, 1:, :] - pred[:, :, :-1, :]
128        target_grad_y = target[:, :, 1:, :] - target[:, :, :-1, :]
129
130        weight_x = weight_map[:, :, :, 1:] * weight_map[:, :, :, :-1]
131        weight_y = weight_map[:, :, 1:, :] * weight_map[:, :, :-1, :]
132
133        loss_x = torch.mean(torch.abs(pred_grad_x - target_grad_x) * weight_x)
134        loss_y = torch.mean(torch.abs(pred_grad_y - target_grad_y) * weight_y)
135        
136        # loss_x = F.l1_loss(pred_grad_x, target_grad_x) 
137        # loss_y = F.l1_loss(pred_grad_y, target_grad_y)
138
139        return loss_x + loss_y
140
141    def ssim_loss(self, pred, target, weight_map):
142        # SSIM returns similarity, so we subtract from 1
143        if pred.ndim == 3:
144            pred = pred.unsqueeze(1)
145        if target.ndim == 3:
146            target = target.unsqueeze(1)
147
148        # self.ssim_module = self.ssim_module.to(pred.device)
149        return self.ssim_module(pred, target)
150
151    def edge_aware_loss(self, pred, target, weight_map):
152        if pred.ndim == 3:
153            pred = pred.unsqueeze(1)
154        if target.ndim == 3:
155            target = target.unsqueeze(1)
156        if weight_map.ndim == 3:
157            weight_map = weight_map.unsqueeze(1)
158
159        pred_grad_x = pred[:, :, :, :-1] - pred[:, :, :, 1:]
160        pred_grad_y = pred[:, :, :-1, :] - pred[:, :, 1:, :]
161
162        target_grad_x = torch.mean(torch.abs(target[:, :, :, :-1] - target[:, :, :, 1:]), 1, keepdim=True)
163        target_grad_y = torch.mean(torch.abs(target[:, :, :-1, :] - target[:, :, 1:, :]), 1, keepdim=True)
164
165        weight_x = weight_map[:, :, :, 1:] * weight_map[:, :, :, :-1]
166        weight_y = weight_map[:, :, 1:, :] * weight_map[:, :, :-1, :]
167
168        pred_grad_x *= torch.exp(-target_grad_x* weight_x) 
169        pred_grad_y *= torch.exp(-target_grad_y* weight_y)
170
171        # return (pred_grad_y.abs().mean() + target_grad_y.abs().mean())
172        return (pred_grad_x.abs().mean() + pred_grad_y.abs().mean())
173
174    def l1_loss(self, pred, target, weight_map):
175        loss = torch.abs(target - pred) * weight_map
176        return loss.mean()
177
178    def variance_loss(self, pred, target):
179        pred_var = torch.var(pred)
180        target_var = torch.var(target)
181        return F.mse_loss(pred_var, target_var)
182    
183    def range_loss(self, pred, target):
184        pred_min, pred_max = torch.min(pred), torch.max(pred)
185        target_min, target_max = torch.min(target), torch.max(target)
186        
187        min_loss = F.mse_loss(pred_min, target_min)
188        max_loss = F.mse_loss(pred_max, target_max)
189        
190        return min_loss + max_loss
191
192    def blur_loss(self, pred, target):
193        laplacian_kernel = torch.tensor([[[[0, 1, 0],
194                                           [1, -4, 1],
195                                           [0, 1, 0]]]], dtype=pred.dtype, device=pred.device)
196
197        if pred.ndim == 3:
198            pred = pred.unsqueeze(1)
199        if target.ndim == 3:
200            target = target.unsqueeze(1)
201
202        pred_lap = F.conv2d(pred, laplacian_kernel, padding=1)
203        target_lap = F.conv2d(target, laplacian_kernel, padding=1)
204
205        return F.l1_loss(pred_lap, target_lap)
206
207    def blur_loss(self, pred, target):
208        laplacian_kernel = torch.tensor([[[[0, 1, 0],
209                                           [1, -4, 1],
210                                           [0, 1, 0]]]], dtype=pred.dtype, device=pred.device)
211
212        if pred.ndim == 3:
213            pred = pred.unsqueeze(1)
214        if target.ndim == 3:
215            target = target.unsqueeze(1)
216
217        pred_lap = F.conv2d(pred, laplacian_kernel, padding=1)
218        target_lap = F.conv2d(target, laplacian_kernel, padding=1)
219
220        return F.l1_loss(pred_lap, target_lap)
221
222    def forward(self, pred, target, weight_map=None, should_calc_weight_map=False):
223        """
224        Computes the weighted combined loss between prediction and target.
225
226        Parameter:
227        - pred (torch.Tensor): 
228            Predicted output tensor.
229        - target (torch.Tensor): 
230            Ground truth tensor.
231        - weight_map (torch.Tensor or None): 
232            Optional pixel-wise weighting map.
233        - should_calc_weight_map (bool): 
234            If True and weight_map is None, calculates a weight map from target.
235
236        Returns:
237        - torch.Tensor: Weighted sum of all losses.
238        """
239        if type(weight_map) == type(None):
240            if should_calc_weight_map:
241                weight_map = calc_weight_map(target)
242            else:
243                # no mask/weight-map
244                # FIXME -> right
245                weight_map = torch.ones_like(pred)
246
247        loss_silog = self.silog_loss(pred, target, weight_map)
248        loss_grad = self.gradient_l1_loss(pred, target, weight_map)
249        loss_ssim = self.ssim_loss(pred, target, weight_map)
250        loss_l1 = self.l1_loss(pred, target, weight_map)
251        loss_edge_aware = self.edge_aware_loss(pred, target, weight_map)
252        loss_var = self.variance_loss(pred, target)
253        loss_range = self.range_loss(pred, target)
254        loss_blur = self.blur_loss(pred, target)
255
256        # reset avgs
257        if self.steps > 24:
258            self.step()
259
260        self.avg_loss_silog += loss_silog
261        self.avg_loss_grad += loss_grad
262        self.avg_loss_ssim += loss_ssim
263        self.avg_loss_l1 += loss_l1
264        self.avg_loss_edge_aware += loss_edge_aware
265        self.avg_loss_var += loss_var
266        self.avg_loss_range += loss_range
267        self.avg_loss_blur += loss_blur
268        self.steps += 1
269
270        total_loss = (
271            self.weight_silog * loss_silog +
272            self.weight_grad * loss_grad +
273            self.weight_ssim * loss_ssim +
274            self.weight_edge_aware * loss_edge_aware +
275            self.weight_l1 * loss_l1 +
276            self.weight_var * loss_var +
277            self.weight_range * loss_range +
278            self.weight_blur * loss_blur
279        )
280
281        return total_loss
282
283    def step(self, epoch=None):
284        """
285        Resets the running averages of all tracked losses.
286        """
287        self.avg_loss_silog = 0
288        self.avg_loss_grad = 0
289        self.avg_loss_ssim = 0
290        self.avg_loss_l1 = 0
291        self.avg_loss_edge_aware = 0
292        self.avg_loss_var = 0
293        self.avg_loss_range = 0
294        self.avg_loss_blur = 0
295        self.steps = 0
296
297    def get_avg_losses(self):
298        """
299        Returns the running average of all individual losses.
300
301        Returns:
302        - tuple: (avg_loss_silog, avg_loss_grad, avg_loss_ssim, avg_loss_l1,
303                avg_loss_edge_aware, avg_loss_var, avg_loss_range, avg_loss_blur)
304        """
305        return (self.avg_loss_silog/self.steps,
306                self.avg_loss_grad/self.steps,
307                self.avg_loss_ssim/self.steps,
308                self.avg_loss_l1/self.steps,
309                self.avg_loss_edge_aware/self.steps,
310                self.avg_loss_var/self.steps,
311                self.avg_loss_range/self.steps,
312                self.avg_loss_blur/self.steps
313               )
314
315    def get_dict(self):
316        """
317        Returns a dictionary of average losses and their corresponding weights.
318
319        Returns:
320        - dict: All loss components with their weights.
321        """
322        loss_silog, loss_grad, loss_ssim, loss_l1, loss_edge_aware, loss_var, loss_range, loss_blur = self.get_avg_losses()
323        return {
324                f"loss_silog": loss_silog, 
325                f"loss_grad": loss_grad, 
326                f"loss_ssim": loss_ssim,
327                f"loss_L1": loss_l1,
328                f"loss_edge aware": loss_edge_aware,
329                f"loss_var": loss_var,
330                f"loss_range": loss_range,
331                f"loss_blur": loss_blur,
332                f"weight_loss_silog": self.weight_silog, 
333                f"weight_loss_grad": self.weight_grad,
334                f"_weight_loss_ssim": self.weight_ssim,
335                f"_weight_loss_L1": self.weight_l1,
336                f"weight_loss_edge_aware": self.weight_edge_aware,
337                f"weight_loss_var": self.weight_var,
338                f"weight_loss_range": self.weight_range,
339                f"weight_loss_blur": self.weight_blur
340               }
341
342def calc_weight_map(target):
343    """
344    Calculates a per-pixel weighting map for a target tensor based on unique value frequencies.
345
346    Less frequent values are given higher weights to emphasize their contribution in loss computations.
347
348    Parameter:
349    - target (torch.Tensor): 
350        Ground truth tensor.
351
352    Returns:
353    - torch.Tensor: Weight map tensor of the same shape as target.
354    """
355    values, counts = torch.unique(target.flatten(), return_counts=True)
356    all_counts = counts.sum().float()
357    
358    # weight_factor = 2.0
359    # weights = {values[idx].item(): max(torch.exp( ( (1-(counts[idx].item()/all_counts))) *weight_factor), 0.0001) for idx in range(len(values))}
360    
361    weights = {values[idx].item(): 255.0/counts[idx].item() for idx in range(len(values))}
362
363    # print(f"Weights:")
364    # for cur_value, cur_counts in list(sorted(weights.items(), key=lambda x:x[0])):
365    #     print('    - '+str(round(cur_value, 4))+': '+str(cur_counts.item()))
366
367    weights_map = torch.zeros_like(target, dtype=torch.float)
368    for cur_value in values:
369        cur_value = cur_value.item()
370        weights_map[target == cur_value] = weights[cur_value]
371
372    return weights_map