<reponame>duanzhiihao/mycv
import torch
import torch.nn as nn

from mycv.external.compressai.zoo import mbt2018


class MBT2018(nn.Module):
    def __init__(self, quality, metric='mse'):
        super().__init__()
        model  = mbt2018(quality, metric=metric, pretrained=True)
        self.g_a = model.g_a
        for p in self.g_a.parameters():
            p.requires_grad_(False)
        self.g_s = model.g_s
        for p in self.g_s.parameters():
            p.requires_grad_(False)

        model  = mbt2018(quality, metric=metric, pretrained=False)
        self.h_a = model.h_a
        self.h_s = model.h_s
        self.entropy_bottleneck = model.entropy_bottleneck
        self.gaussian_conditional = model.gaussian_conditional
        self.context_prediction = model.context_prediction
        self.entropy_parameters = model.entropy_parameters

    def forward(self, x):
        y = self.g_a(x)
        z = self.h_a(y)
        z_hat, z_likelihoods = self.entropy_bottleneck(z)
        params = self.h_s(z_hat)

        y_hat = self.gaussian_conditional.quantize(
            y, "noise" if self.training else "dequantize"
        )
        ctx_params = self.context_prediction(y_hat)
        gaussian_params = self.entropy_parameters(
            torch.cat((params, ctx_params), dim=1)
        )
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        _, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)
        x_hat = self.g_s(y_hat)

        return {
            "x_hat": x_hat,
            "likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
        }

    def forward_entropy(self, x: torch.Tensor):
        assert not x.requires_grad
        y = self.g_a(x)
        assert not y.requires_grad

        z = self.h_a(y)
        z_hat, z_likelihoods = self.entropy_bottleneck(z)
        params = self.h_s(z_hat)

        y_hat = self.gaussian_conditional.quantize(
            y, "noise" if self.training else "dequantize"
        )
        ctx_params = self.context_prediction(y_hat)
        gaussian_params = self.entropy_parameters(
            torch.cat((params, ctx_params), dim=1)
        )
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        _, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)
        # x_hat = self.g_s(y_hat)

        return {
            # "x_hat": x_hat,
            "likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
        }


class CompressedVAE(nn.Module):
    def __init__(self, channels=192):
        super().__init__()

    def forward(self, x):
        return x


class MyNIC(nn.Module):
    def __init__(self, quality, metric='mse'):
        super().__init__()
        model  = mbt2018(quality, metric=metric, pretrained=True)
        self.g_a = model.g_a
        for p in self.g_a.parameters():
            p.requires_grad_(False)
        self.g_s = model.g_s
        for p in self.g_s.parameters():
            p.requires_grad_(False)

        from mycv.models.probabilistic.latentvae import VDVAE
        self.entropy = VDVAE('test')

    def forward(self, x):
        z = self.g_a(x)
        
        stats = self.entropy(z)
        imH, imW = x.shape[2:4]
        zC, zH, zW = z.shape[1:4]
        bpp = stats['bpdim'] * (zC*zH*zW/imH/imW)
        zhat = stats['zhat']

        x_hat = self.g_s(zhat)

        return {
            "x_hat": x_hat,
            "bpp": bpp,
        }

    def forward_entropy(self, x: torch.Tensor):
        assert not x.requires_grad
        z = self.g_a(x)
        assert not z.requires_grad

        stats = self.entropy(z)
        imH, imW = x.shape[2:4]
        zC, zH, zW = z.shape[1:4]
        bpp = stats['bpdim'] * (zC*zH*zW/imH/imW)

        return {
            "bpp": bpp,
        }

    @torch.no_grad()
    def uncond_sample(self, bhw=(4,256,256)):
        nB, nH, nW = bhw
        z_samples = self.entropy.uncond_sample(bhw=(nB,nH//16,nW//16))
        samples = self.g_s(z_samples)
        return samples


def main():
    debug = 1
    import matplotlib.pyplot as plt
    import mycv.utils.image as imu
    from mycv.datasets.imcoding import nic_evaluate
    # model = cheng2020_anchor(quality=4, metric='ms-ssim', pretrained=True)
    # model = MBT2018(quality=4, metric='mse')
    model = MyNIC(quality=4, metric='mse')
    model = model.cuda()
    model.eval()

    checkpoint = torch.load('mycv/weights/latentvae.pt')
    model.load_state_dict(checkpoint['model'])

    # results = nic_evaluate(model)
    # print(results); exit()

    samples = model.uncond_sample(bhw=(4,256,256))
    for sample in samples:
        sampled = imu.to_numpy_img(sample)
        plt.figure(); plt.imshow(sampled)
    plt.show()
    debug = 1


if __name__ == '__main__':
    main()

import torch
import torch.nn as nn

from mycv.external.compressai.zoo import mbt2018


class MBT2018(nn.Module):
    def __init__(self, quality, metric='mse'):
        super().__init__()
        model  = mbt2018(quality, metric=metric, pretrained=True)
        self.g_a = model.g_a
        for p in self.g_a.parameters():
            p.requires_grad_(False)
        self.g_s = model.g_s
        for p in self.g_s.parameters():
            p.requires_grad_(False)

        model  = mbt2018(quality, metric=metric, pretrained=False)
        self.h_a = model.h_a
        self.h_s = model.h_s
        self.entropy_bottleneck = model.entropy_bottleneck
        self.gaussian_conditional = model.gaussian_conditional
        self.context_prediction = model.context_prediction
        self.entropy_parameters = model.entropy_parameters

    def forward(self, x):
        y = self.g_a(x)
        z = self.h_a(y)
        z_hat, z_likelihoods = self.entropy_bottleneck(z)
        params = self.h_s(z_hat)

        y_hat = self.gaussian_conditional.quantize(
            y, "noise" if self.training else "dequantize"
        )
        ctx_params = self.context_prediction(y_hat)
        gaussian_params = self.entropy_parameters(
            torch.cat((params, ctx_params), dim=1)
        )
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        _, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)
        x_hat = self.g_s(y_hat)

        return {
            "x_hat": x_hat,
            "likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
        }

    def forward_entropy(self, x: torch.Tensor):
        assert not x.requires_grad
        y = self.g_a(x)
        assert not y.requires_grad

        z = self.h_a(y)
        z_hat, z_likelihoods = self.entropy_bottleneck(z)
        params = self.h_s(z_hat)

        y_hat = self.gaussian_conditional.quantize(
            y, "noise" if self.training else "dequantize"
        )
        ctx_params = self.context_prediction(y_hat)
        gaussian_params = self.entropy_parameters(
            torch.cat((params, ctx_params), dim=1)
        )
        scales_hat, means_hat = gaussian_params.chunk(2, 1)
        _, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat)
        # x_hat = self.g_s(y_hat)

        return {
            # "x_hat": x_hat,
            "likelihoods": {"y": y_likelihoods, "z": z_likelihoods},
        }


class CompressedVAE(nn.Module):
    def __init__(self, channels=192):
        super().__init__()

    def forward(self, x):
        return x


class MyNIC(nn.Module):
    def __init__(self, quality, metric='mse'):
        super().__init__()
        model  = mbt2018(quality, metric=metric, pretrained=True)
        self.g_a = model.g_a
        for p in self.g_a.parameters():
            p.requires_grad_(False)
        self.g_s = model.g_s
        for p in self.g_s.parameters():
            p.requires_grad_(False)

        from mycv.models.probabilistic.latentvae import VDVAE
        self.entropy = VDVAE('test')

    def forward(self, x):
        z = self.g_a(x)
        
        stats = self.entropy(z)
        imH, imW = x.shape[2:4]
        zC, zH, zW = z.shape[1:4]
        bpp = stats['bpdim'] * (zC*zH*zW/imH/imW)
        zhat = stats['zhat']

        x_hat = self.g_s(zhat)

        return {
            "x_hat": x_hat,
            "bpp": bpp,
        }

    def forward_entropy(self, x: torch.Tensor):
        assert not x.requires_grad
        z = self.g_a(x)
        assert not z.requires_grad

        stats = self.entropy(z)
        imH, imW = x.shape[2:4]
        zC, zH, zW = z.shape[1:4]
        bpp = stats['bpdim'] * (zC*zH*zW/imH/imW)

        return {
            "bpp": bpp,
        }

    @torch.no_grad()
    def uncond_sample(self, bhw=(4,256,256)):
        nB, nH, nW = bhw
        z_samples = self.entropy.uncond_sample(bhw=(nB,nH//16,nW//16))
        samples = self.g_s(z_samples)
        return samples


def main():
    debug = 1
    import matplotlib.pyplot as plt
    import mycv.utils.image as imu
    from mycv.datasets.imcoding import nic_evaluate
    # model = cheng2020_anchor(quality=4, metric='ms-ssim', pretrained=True)
    # model = MBT2018(quality=4, metric='mse')
    model = MyNIC(quality=4, metric='mse')
    model = model.cuda()
    model.eval()

    checkpoint = torch.load('mycv/weights/latentvae.pt')
    model.load_state_dict(checkpoint['model'])

    # results = nic_evaluate(model)
    # print(results); exit()

    samples = model.uncond_sample(bhw=(4,256,256))
    for sample in samples:
        sampled = imu.to_numpy_img(sample)
        plt.figure(); plt.imshow(sampled)
    plt.show()
    debug = 1


if __name__ == '__main__':
    main()
# x_hat = self.g_s(y_hat) # "x_hat": x_hat, # model = cheng2020_anchor(quality=4, metric='ms-ssim', pretrained=True) # model = MBT2018(quality=4, metric='mse') # results = nic_evaluate(model) # print(results); exit()