FGT_codes/FGT/models/utils/loss.py

import os
import sys

import torch
import torch.nn as nn
import numpy as np


class AlignLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.loss_fn = nn.L1Loss(reduction=reduction)

    def forward(self, frames, masks, aligned_vs, aligned_rs):
        """

        :param frames: The original frames(GT)
        :param masks: Original masks
        :param aligned_vs: aligned visibility map from reference frame(List: B, C, T, H, W)
        :param aligned_rs: aligned reference frames(List: B, C, T, H, W)
        :return:
        """
        try:
            B, C, T, H, W = frames.shape
        except ValueError:
            frames = frames.unsqueeze(2)
            masks = masks.unsqueeze(2)
        B, C, T, H, W = frames.shape
        loss = 0
        for i in range(T):
            frame = frames[:, :, i]
            mask = masks[:, :, i]
            aligned_v = aligned_vs[i]
            aligned_r = aligned_rs[i]
            loss += self._singleFrameAlignLoss(frame, mask, aligned_v, aligned_r)
        return loss

    def _singleFrameAlignLoss(self, targetFrame, targetMask, aligned_v, aligned_r):
        """

        :param targetFrame: targetFrame to be aligned-> B, C, H, W
        :param targetMask: the mask of target frames
        :param aligned_v: aligned visibility map from reference frame
        :param aligned_r: aligned reference frame-> B, C, T, H, W
        :return:
        """
        targetVisibility = 1. - targetMask
        targetVisibility = targetVisibility.unsqueeze(2)
        targetFrame = targetFrame.unsqueeze(2)
        visibility_map = targetVisibility * aligned_v
        target_visibility = visibility_map * targetFrame
        reference_visibility = visibility_map * aligned_r
        loss = 0
        for i in range(aligned_r.shape[2]):
            loss += self.loss_fn(target_visibility[:, :, i], reference_visibility[:, :, i])
        return loss


class HoleVisibleLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.loss_fn = nn.L1Loss(reduction=reduction)

    def forward(self, outputs, masks, GTs, c_masks):
        try:
            B, C, T, H, W = outputs.shape
        except ValueError:
            outputs = outputs.unsqueeze(2)
            masks = masks.unsqueeze(2)
            GTs = GTs.unsqueeze(2)
            c_masks = c_masks.unsqueeze(2)
        B, C, T, H, W = outputs.shape
        loss = 0
        for i in range(T):
            loss += self._singleFrameHoleVisibleLoss(outputs[:, :, i], masks[:, :, i], c_masks[:, :, i], GTs[:, :, i])
        return loss

    def _singleFrameHoleVisibleLoss(self, targetFrame, targetMask, c_mask, GT):
        return self.loss_fn(targetMask * c_mask * targetFrame, targetMask * c_mask * GT)


class HoleInvisibleLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.loss_fn = nn.L1Loss(reduction=reduction)

    def forward(self, outputs, masks, GTs, c_masks):
        try:
            B, C, T, H, W = outputs.shape
        except ValueError:
            outputs = outputs.unsqueeze(2)
            masks = masks.unsqueeze(2)
            GTs = GTs.unsqueeze(2)
            c_masks = c_masks.unsqueeze(2)
        B, C, T, H, W = outputs.shape
        loss = 0
        for i in range(T):
            loss += self._singleFrameHoleInvisibleLoss(outputs[:, :, i], masks[:, :, i], c_masks[:, :, i], GTs[:, :, i])
        return loss

    def _singleFrameHoleInvisibleLoss(self, targetFrame, targetMask, c_mask, GT):
        return self.loss_fn(targetMask * (1. - c_mask) * targetFrame, targetMask * (1. - c_mask) * GT)


class NonHoleLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.loss_fn = nn.L1Loss(reduction=reduction)

    def forward(self, outputs, masks, GTs):
        try:
            B, C, T, H, W = outputs.shape
        except ValueError:
            outputs = outputs.unsqueeze(2)
            masks = masks.unsqueeze(2)
            GTs = GTs.unsqueeze(2)
        B, C, T, H, W = outputs.shape
        loss = 0
        for i in range(T):
            loss += self._singleNonHoleLoss(outputs[:, :, i], masks[:, :, i], GTs[:, :, i])
        return loss

    def _singleNonHoleLoss(self, targetFrame, targetMask, GT):
        return self.loss_fn((1. - targetMask) * targetFrame, (1. - targetMask) * GT)


class ReconLoss(nn.Module):
    def __init__(self, reduction='mean', masked=False):
        super().__init__()
        self.loss_fn = nn.L1Loss(reduction=reduction)
        self.masked = masked

    def forward(self, model_output, target, mask):
        outputs = model_output
        targets = target
        if self.masked:
            masks = mask
            return self.loss_fn(outputs * masks, targets * masks)  # L1 loss in masked region
        else:
            return self.loss_fn(outputs, targets)  # L1 loss in the whole region


class VGGLoss(nn.Module):
    def __init__(self, vgg):
        super().__init__()
        self.l1_loss = nn.L1Loss()
        self.vgg = vgg

    def vgg_loss(self, output, target):
        output_feature = self.vgg(output)
        target_feature = self.vgg(target)
        loss = (
                self.l1_loss(output_feature.relu2_2, target_feature.relu2_2)
                + self.l1_loss(output_feature.relu3_3, target_feature.relu3_3)
                + self.l1_loss(output_feature.relu4_3, target_feature.relu4_3)
        )
        return loss

    def forward(self, data_input, model_output):
        targets = data_input
        outputs = model_output
        mean_image_loss = self.vgg_loss(outputs, targets)
        return mean_image_loss


class StyleLoss(nn.Module):
    def __init__(self, vgg, original_channel_norm=True):
        super().__init__()
        self.l1_loss = nn.L1Loss()
        self.vgg = vgg
        self.original_channel_norm = original_channel_norm

    # From https://github.com/pytorch/tutorials/blob/master/advanced_source/neural_style_tutorial.py
    def gram_matrix(self, input):
        a, b, c, d = input.size()  # a=batch size(=1)
        # b=number of feature maps
        # (c,d)=dimensions of a f. map (N=c*d)

        features = input.view(a * b, c * d)  # resise F_XL into \hat F_XL

        G = torch.mm(features, features.t())  # compute the gram product

        # we 'normalize' the values of the gram matrix
        # by dividing by the number of element in each feature maps.
        return G.div(a * b * c * d)

    # Implement "Image Inpainting for Irregular Holes Using Partial Convolutions", Liu et al., 2018
    def style_loss(self, output, target):
        output_features = self.vgg(output)
        target_features = self.vgg(target)
        layers = ['relu2_2', 'relu3_3', 'relu4_3']  # n_channel: 128 (=2 ** 7), 256 (=2 ** 8), 512 (=2 ** 9)
        loss = 0
        for i, layer in enumerate(layers):
            output_feature = getattr(output_features, layer)
            target_feature = getattr(target_features, layer)
            B, C_P, H, W = output_feature.shape
            output_gram_matrix = self.gram_matrix(output_feature)
            target_gram_matrix = self.gram_matrix(target_feature)
            if self.original_channel_norm:
                C_P_square_divider = 2 ** (i + 1)  # original design (avoid too small loss)
            else:
                C_P_square_divider = C_P ** 2
                assert C_P == 128 * 2 ** i
            loss += self.l1_loss(output_gram_matrix, target_gram_matrix) / C_P_square_divider
        return loss

    def forward(self, data_input, model_output):
        targets = data_input
        outputs = model_output
        mean_image_loss = self.style_loss(outputs, targets)
        return mean_image_loss


class L1LossMaskedMean(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.L1Loss(reduction='sum')

    def forward(self, x, y, mask):
        masked = 1 - mask  # 默认missing region的mask值为0，原有区域为1
        l1_sum = self.l1(x * masked, y * masked)
        return l1_sum / torch.sum(masked)


class L2LossMaskedMean(nn.Module):
    def __init__(self, reduction='sum'):
        super().__init__()
        self.l2 = nn.MSELoss(reduction=reduction)

    def forward(self, x, y, mask):
        masked = 1 - mask
        l2_sum = self.l2(x * masked, y * masked)
        return l2_sum / torch.sum(masked)


class ImcompleteVideoReconLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss_fn = L1LossMaskedMean()

    def forward(self, data_input, model_output):
        imcomplete_video = model_output['imcomplete_video']
        targets = data_input['targets']
        down_sampled_targets = nn.functional.interpolate(
            targets.transpose(1, 2), scale_factor=[1, 0.5, 0.5])

        masks = data_input['masks']
        down_sampled_masks = nn.functional.interpolate(
            masks.transpose(1, 2), scale_factor=[1, 0.5, 0.5])
        return self.loss_fn(
            imcomplete_video, down_sampled_targets,
            down_sampled_masks
        )


class CompleteFramesReconLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.loss_fn = L1LossMaskedMean()

    def forward(self, data_input, model_output):
        outputs = model_output['outputs']
        targets = data_input['targets']
        masks = data_input['masks']
        return self.loss_fn(outputs, targets, masks)


class AdversarialLoss(nn.Module):
    r"""
    Adversarial loss
    https://arxiv.org/abs/1711.10337
    """

    def __init__(self, type='nsgan', target_real_label=1.0, target_fake_label=0.0):
        r"""
        type = nsgan | lsgan | hinge
        """
        super(AdversarialLoss, self).__init__()
        self.type = type
        self.register_buffer('real_label', torch.tensor(target_real_label))
        self.register_buffer('fake_label', torch.tensor(target_fake_label))

        if type == 'nsgan':
            self.criterion = nn.BCELoss()
        elif type == 'lsgan':
            self.criterion = nn.MSELoss()
        elif type == 'hinge':
            self.criterion = nn.ReLU()

    def __call__(self, outputs, is_real, is_disc=None):
        if self.type == 'hinge':
            if is_disc:
                if is_real:
                    outputs = -outputs
                return self.criterion(1 + outputs).mean()
            else:
                return (-outputs).mean()
        else:
            labels = (self.real_label if is_real else self.fake_label).expand_as(
                outputs)
            loss = self.criterion(outputs, labels)
            return loss


# # From https://github.com/phoenix104104/fast_blind_video_consistency
# class TemporalWarpingLoss(nn.Module):
#     def __init__(self, opts, flownet_checkpoint_path=None, alpha=50):
#         super().__init__()
#         self.loss_fn = L1LossMaskedMean()
#         self.alpha = alpha
#         self.opts = opts
#
#         assert flownet_checkpoint_path is not None, "Flownet2 pretrained models must be provided"
#
#         self.flownet_checkpoint_path = flownet_checkpoint_path
#         raise NotImplementedError
#
#     def get_flownet_checkpoint_path(self):
#         return self.flownet_checkpoint_path
#
#     def _flownetwrapper(self):
#         Flownet = FlowNet2(self.opts, requires_grad=False)
#         Flownet2_ckpt = torch.load(self.flownet_checkpoint_path)
#         Flownet.load_state_dict(Flownet2_ckpt['state_dict'])
#         Flownet.to(device)
#         Flownet.exal()
#         return Flownet
#
#     def _setup(self):
#         self.flownet = self._flownetwrapper()
#
#     def _get_non_occlusuib_mask(self, targets, warped_targets):
#         non_occlusion_masks = torch.exp(
#             -self.alpha * torch.sum(targets[:, 1:] - warped_targets, dim=2).pow(2)
#         ).unsqueeze(2)
#         return non_occlusion_masks
#
#     def _get_loss(self, outputs, warped_outputs, non_occlusion_masks, masks):
#         return self.loss_fn(
#             outputs[:, 1:] * non_occlusion_masks,
#             warped_outputs * non_occlusion_masks,
#             masks[:, 1:]
#         )
#
#     def forward(self, data_input, model_output):
#         if self.flownet is None:
#             self._setup()
#
#         targets = data_input['targets'].to(device)
#         outputs = model_output['outputs'].to(device)
#         flows = self.flownet.infer_video(targets).to(device)
#
#         from utils.flow_utils import warp_optical_flow
#         warped_targets = warp_optical_flow(targets[:, :-1], -flows).detach()
#         warped_outputs = warp_optical_flow(outputs[:, :-1], -flows).detach()
#         non_occlusion_masks = self._get_non_occlusion_mask(targets, warped_targets)
#
#         # model_output is passed by name and dictionary is mutable
#         # These values are sent to trainer for visualization
#         model_output['warped_outputs'] = warped_outputs[0]
#         model_output['warped_targets'] = warped_targets[0]
#         model_output['non_occlusion_masks'] = non_occlusion_masks[0]
#         from utils.flow_utils import flow_to_image
#         flow_imgs = []
#         for flow in flows[0]:
#             flow_img = flow_to_image(flow.cpu().permute(1, 2, 0).detach().numpy()).transpose(2, 0, 1)
#             flow_imgs.append(torch.Tensor(flow_img))
#         model_output['flow_imgs'] = flow_imgs
#
#         masks = data_input['masks'].to(device)
#         return self._get_loss(outputs, warped_outputs, non_occlusion_masks, masks)
#
#
# class TemporalWarpingError(TemporalWarpingLoss):
#     def __init__(self, flownet_checkpoint_path, alpha=50):
#         super().__init__(flownet_checkpoint_path, alpha)
#         self.loss_fn = L2LossMaskedMean(reduction='none')
#
#     def _get_loss(self, outputs, warped_outputs, non_occlusion_masks, masks):
#         # See https://arxiv.org/pdf/1808.00449.pdf 4.3
#         # The sum of non_occlusion_masks is different for each video,
#         # So the batch dim is kept
#         loss = self.loss_fn(
#             outputs[:, 1:] * non_occlusion_masks,
#             warped_outputs * non_occlusion_masks,
#             masks[:, 1:]
#         ).sum(1).sum(1).sum(1).sum(1)
#
#         loss = loss / non_occlusion_masks.sum(1).sum(1).sum(1).sum(1)
#         return loss.sum()


class ValidLoss(nn.Module):
    def __init__(self):
        super(ValidLoss, self).__init__()
        self.loss_fn = nn.L1Loss(reduction='mean')

    def forward(self, model_output, target, mk):
        outputs = model_output
        targets = target
        return self.loss_fn(outputs * (1 - mk), targets * (1 - mk))  # L1 loss in masked region



class TVLoss(nn.Module):
    def __init__(self):
        super(TVLoss, self).__init__()

    def forward(self, mask_input, model_output):
        # View 3D data as 2D
        outputs = model_output

        if len(mask_input.shape) == 4:
            mask_input = mask_input.unsqueeze(2)
        if len(outputs.shape) == 4:
            outputs = outputs.unsqueeze(2)

        outputs = outputs.permute((0, 2, 1, 3, 4)).contiguous()
        masks = mask_input.permute((0, 2, 1, 3, 4)).contiguous()

        B, L, C, H, W = outputs.shape
        x = outputs.view([B * L, C, H, W])

        masks = masks.view([B * L, -1])
        mask_areas = masks.sum(dim=1)

        h_x = x.size()[2]
        w_x = x.size()[3]
        h_tv = torch.pow((x[:, :, 1:, :] - x[:, :, :h_x - 1, :]), 2).sum(1).sum(1).sum(1)  # 差分是为了求梯度，本质上还是梯度平方和
        w_tv = torch.pow((x[:, :, :, 1:] - x[:, :, :, :w_x - 1]), 2).sum(1).sum(1).sum(1)
        return ((h_tv + w_tv) / mask_areas).mean()


# for debug
def show_images(image, name):
    import cv2
    import numpy as np
    image = np.array(image)
    image[image > 0.5] = 255.
    image = image.transpose((1, 2, 0))
    cv2.imwrite(name, image)


if __name__ == '__main__':
    # test align loss,
    targetFrame = torch.ones(1, 3, 32, 32)
    GT = torch.ones(1, 3, 32, 32)
    GT += 1
    mask = torch.zeros(1, 1, 32, 32)
    mask[:, :, 8:24, 8:24] = 1.

    # referenceFrames = torch.ones(1, 3, 4, 32, 32)
    # referenceMasks = torch.zeros(1, 1, 4, 32, 32)
    # referenceMasks[:, :, 0, 4:12, 4:12] = 1.
    # referenceFrames[:, :, 0, 4:12, 4:12] = 2.
    # referenceMasks[:, :, 1, 4:12, 20:28] = 1.
    # referenceFrames[:, :, 1, 4:12, 20:28] = 2.
    # referenceMasks[:, :, 2, 20:28, 4:12] = 1.
    # referenceFrames[:, :, 2, 20:28, 4:12] = 2.
    # referenceMasks[:, :, 3, 20:28, 20:28] = 1.
    # referenceFrames[:, :, 3, 20:28, 20:28] = 2.
    #
    # aligned_v = referenceMasks
    # aligned_v, referenceFrames = [aligned_v], [referenceFrames]
    #
    # result = AlignLoss()(targetFrame, mask, aligned_v, referenceFrames)
    # print(result)

    c_mask = torch.zeros(1, 1, 32, 32)
    c_mask[:, :, 8:16, 16:24] = 1.
    result1 = HoleVisibleLoss()(targetFrame, mask, GT, c_mask)
    result2 = HoleInvisibleLoss()(targetFrame, mask, GT, c_mask)
    result3 = NonHoleLoss()(targetFrame, mask, GT)
    print('vis: {}, invis: {}, gt: {}'.format(result1, result2, result3))