BatchTest.py

# -*- coding: utf-8 -*-
"""
Created on Mon Apr  4 10:55:54 2022

@author: Mateo

plotqf
"""

import os.path
import numpy as np
from PIL import Image
import io
from piq import psnr, ssim
from piq import multi_scale_ssim as ms_ssim
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
import torchvision.transforms as tt
from io import BytesIO
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import bz2
import torch.nn.functional as F
import pandas as pd
from torchvision.utils import make_grid, save_image

#Paths
dspath = "E:/MachineLearning/datasets/"
structpath = "E:/MachineLearning/RinRUnpix/finalOptions/"

#Prepare test data
torch.set_num_threads(4)
valid_tfms = tt.Compose([ tt.ToTensor() ])
valid_ds = ImageFolder(dspath + 'kodak', valid_tfms)
valid_dl = DataLoader(valid_ds, 1)
numc = 4
global counter 
counter = 0

results = []

structdir = structpath + str(numc) + '/UtMq-cprQ2-clipint8-to0-fmod'
loaddir = structdir + '/autoenc00115.pth'
sdirect = structdir + '/test/KODAKOUT/'    

if torch.cuda.is_available():  
    device = "cuda:0" 
else:  
    device = "cpu"  
print(device)

# to save the images generated by the generator
def save_generator_image(image, path):
    save_image(image, path)

def tobytes(latint):
    np_bytes = BytesIO() #byte transformer
    np.save(np_bytes, latint, allow_pickle=True) #transform to bytes
    np_bytes = np_bytes.getvalue() #get bytes
    return np_bytes
    
def frombytes(np_bytes):
    load_bytes = BytesIO(np_bytes) #byte loader
    loaded_np = np.load(load_bytes, allow_pickle=True) #get np array
    return loaded_np

def appendbytes(bpady, bpadx, np_bytes):
    #header = bpady[:-1]
    np_bytes = np_bytes + bpady[-2:] + bpadx[-1:]
    return np_bytes
    
def popbytes(rec_latnpby):
    bpady, bpadx = int.from_bytes(rec_latnpby[-2:-1], 'little'), int.from_bytes(rec_latnpby[-1:], 'little')
    rec_latnpby = rec_latnpby[:-3]
    return rec_latnpby, bpady, bpadx

def getprobs(latint, maxmin):
    flat = latint.flatten()
    if maxmin != 0:
        freq,_,_ = plt.hist(flat, maxmin)
    else:
        freq,_,_ = plt.hist(flat)
        
    
    return freq/flat.size

def tonumpy(latent):
    latnp = latent.to('cpu').detach().numpy()
    return latnp.squeeze()

def fromnumpy(rec_latnp):
    rec_unqlatt = torch.from_numpy(np.expand_dims(rec_latnp,0)).to(device)
    return rec_unqlatt

def toint(latnp):
    return np.around(np.around(latnp,2)*100,0).astype(np.int8)
    
def tofloat(latint):
    return (latint.astype(np.float32)/100)

def tohalf(latint):
    if latint.max() < 5 and latint.min() > -6: #[-5;4] = -5-4-3-2-101234 + 5 -> [0;9]
        latint = latint+5
        l1h, l2h = np.split(latint,2) #ex: 32 l1h-> [0:15] l2h-> [16:31]
        #latintj = 
    elif latint.max() < 6 and latint.min() > -5: #[-4;5]
        latint = latint+5
        l1h, l2h = np.split(latint,2) #ex: 32 l1h-> [0:15] l2h-> [16:31]
        #latintj = 
    if  l1h.min() >=0 and l2h.min() >= 0:
        latint = np.core.defchararray.add(l1h.astype(str),l2h.astype(str)).astype(int)
    return latint
        
def fromhalf(latint):
    ##if latint.min() >= 0:
    ##    l1h, l2h = latint
    pass

def savefile(latnpby):
        f = bz2.BZ2File(sdirect + 'latby_bz2{0:05d}.bz2'.format(counter), "wb", compresslevel=9)
        f.write(latnpby)
        f.close()

def readfile():
    f = bz2.BZ2File(sdirect + 'latby_bz2{0:05d}.bz2'.format(counter), 'r')
    bytesnp = f.read()
    f.close()
    return bytesnp
    loaded_np = frombytes(bytesnp)
    #DECOMPRESS
    ##########################
    rec_unqlat = tofloat(loaded_np)
    rec_unqlatt = torch.from_numpy(np.expand_dims(rec_unqlat,0))
    #print((rec_unqlat==latnp).all(), (rec_unqlat==np.around(latnp,2)).all())
    rec_unqlatt = rec_unqlatt.to(device)
    return rec_unqlatt

def imgpad(imgt):
    _,_, yy, xx = imgt.size()
    ogyy, ogxx = yy, xx
    while True:
        if yy%8 != 0:
            imgt = F.pad(imgt, (0,0,0,1), 'reflect') #pad 1 in one side to y
        if xx%8 != 0:
            imgt = F.pad(imgt, (0,1,0,0), 'reflect') #pad 1 in one side to x
        
        _,_,yy,xx = imgt.size()
        if yy%8 ==0 and xx%8==0:
            break
    _,_,pyy,pxx = imgt.size()
    pady = np.int8(pyy-ogyy)
    padx = np.int8(pxx-ogxx)
    bpady, bpadx = tobytes(pady), tobytes(padx)
    return imgt, bpady, bpadx

class STEFunction(torch.autograd.Function):

    @staticmethod
    def forward(ctx, input):
        lquant = torch.quantize_per_tensor(input, 0.01, 0, dtype=torch.qint8) #3 8 decimals
        unqlat = torch.dequantize(lquant)
        return unqlat

    @staticmethod
    def backward(ctx, grad_output):
        #return F.hardtanh(grad_output)
        return grad_output

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

    def forward(self, x):
        x = STEFunction.apply(x)
        return x

class ResidualDenseBlock_5C(nn.Module):
    def __init__(self, nf=64, gc=32, bias=True):
        super(ResidualDenseBlock_5C, self).__init__()
        # gc: growth channel, i.e. intermediate channels
        self.conv1 = nn.Conv2d(nf, gc, 3, 1, 1, bias=bias)
        self.conv2 = nn.Conv2d(nf + gc, gc, 3, 1, 1, bias=bias)
        self.conv3 = nn.Conv2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
        self.conv4 = nn.Conv2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
        self.conv5 = nn.Conv2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        return x5 * 0.2 + x

class RRDB(nn.Module):
    '''Residual in Residual Dense Block'''

    def __init__(self, nf, gc=32):
        super(RRDB, self).__init__()
        self.RDB1 = ResidualDenseBlock_5C(nf, gc)
        self.RDB2 = ResidualDenseBlock_5C(nf, gc)
        self.RDB3 = ResidualDenseBlock_5C(nf, gc)

    def forward(self, x):
        out = self.RDB1(x)
        out = self.RDB2(out)
        out = self.RDB3(out)
        return out * 0.2 + x

class ResidualDenseBlock_5C_dec(nn.Module):
    def __init__(self, nf=64, gc=32, bias=True):
        super(ResidualDenseBlock_5C_dec, self).__init__()
        # gc: growth channel, i.e. intermediate channels
        self.conv1 = nn.ConvTranspose2d(nf, gc, 3, 1, 1, bias=bias)
        self.conv2 = nn.ConvTranspose2d(nf + gc, gc, 3, 1, 1, bias=bias)
        self.conv3 = nn.ConvTranspose2d(nf + 2 * gc, gc, 3, 1, 1, bias=bias)
        self.conv4 = nn.ConvTranspose2d(nf + 3 * gc, gc, 3, 1, 1, bias=bias)
        self.conv5 = nn.ConvTranspose2d(nf + 4 * gc, nf, 3, 1, 1, bias=bias)
        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)

        # initialization
        # mutil.initialize_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1)

    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        return x5 * 0.2 + x

class RRDB_dec(nn.Module):
    '''Residual in Residual Dense Block'''

    def __init__(self, nf, gc=32):
        super(RRDB_dec, self).__init__()
        self.RDB1 = ResidualDenseBlock_5C_dec(nf, gc)
        self.RDB2 = ResidualDenseBlock_5C_dec(nf, gc)
        self.RDB3 = ResidualDenseBlock_5C_dec(nf, gc)

    def forward(self, x):
        out = self.RDB1(x)
        out = self.RDB2(out)
        out = self.RDB3(out)
        return out * 0.2 + x

def conv(ni, nf, ks=3, stride=1, padding=1, **kwargs):
    _conv = nn.Conv2d(ni, nf, kernel_size=ks,stride=stride,padding=padding, **kwargs)
    nn.init.kaiming_normal_(_conv.weight, mode='fan_out')
    return _conv

class resConv(nn.Module):
    def __init__(self):
        super(resConv, self).__init__()
        
        self.enc1 = nn.Sequential(conv(3, 16, 3, 1, 1,padding_mode='reflect'),
                                  nn.LeakyReLU(inplace=True),
                                  nn.PixelUnshuffle(2),
                                  )
        self.enc2 = nn.Sequential(nn.Conv2d(64, 32, 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  )
        self.RinR = RRDB(nf=128, gc=256)

        #latent
        self.enc3 = nn.Sequential(nn.Conv2d(128, int(numc/4), 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  )

        self.lat = StraightThroughEstimator()
        
        #DEC
        self.px1 = nn.Sequential(nn.Conv2d(numc, 512, 3, stride=1, padding=1, padding_mode='reflect'), 
                                 nn.PixelShuffle(2)
                                 )
        
        self.RinRdec = RRDB_dec(nf=128, gc=256)

        self.px2 = nn.Sequential(conv(128, 256, 3, stride=1, padding=1,padding_mode='reflect'), 
                                 nn.PixelShuffle(2),
                                 nn.LeakyReLU(inplace=True)
                                 )
        
        
        self.px3 = nn.Sequential(conv(64, 12, 3, stride=1, padding=1,padding_mode='reflect'),
                                 nn.PixelShuffle(2),
                                 )
        
        self.ecpr = nn.Sequential(nn.Conv2d(3, 16, 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  nn.PixelUnshuffle(2),
                                  nn.PixelUnshuffle(2),
                                  nn.Conv2d(1024, numc, 3, 1, 1,padding_mode='reflect'),
                                  )
        self.dcpr = nn.Sequential(nn.Conv2d(numc, 192, 3, stride=1, padding=1, padding_mode='reflect'), 
                                  nn.PixelShuffle(2),
                                  nn.PixelShuffle(2),
                                  nn.PixelShuffle(2),
                                  )
      
    def encoder(self,x):
        oute = self.enc1(x)     
        x = self.ecpr(x)
        oute = self.enc2(oute)
        oute = self.RinR(oute)
        latent = self.lat(self.enc3(oute) + 0.2*x)
        return latent
    
    def decoder(self,unqlat):
        outd = self.px1(unqlat)
        unqlat = self.dcpr(unqlat)
        outd = self.RinRdec(outd)
        outd = self.px2(outd)
        out = self.px3(outd) + 0.2*unqlat
        out = out.clamp(0,1)
        return out
        
    def forward(self,x):
        latent = self.encoder(x)
        out = self.decoder(latent)
        return out, latent

## ############################################################################
class resConvfmod(nn.Module):
    def __init__(self):
        super(resConvfmod, self).__init__()
        
        self.enc1 = nn.Sequential(conv(3, 16, 3, 1, 1,padding_mode='reflect'),
                                  nn.LeakyReLU(),
                                  nn.PixelUnshuffle(2),
                                  )
        self.enc2 = nn.Sequential(nn.Conv2d(64, 32, 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  )
        self.RinR = RRDB(nf=128, gc=256)

        #latent
        self.enc3 = nn.Sequential(nn.Conv2d(128, int(numc/4), 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  )

        self.lat = StraightThroughEstimator()
        
        #DEC
        self.px1 = nn.Sequential(nn.Conv2d(numc, 512, 3, stride=1, padding=1, padding_mode='reflect'), 
                                 nn.PixelShuffle(2)
                                 )
        
        self.RinRdec = RRDB_dec(nf=128, gc=256)

        self.px2 = nn.Sequential(conv(128, 256, 3, stride=1, padding=1,padding_mode='reflect'), 
                                 nn.PixelShuffle(2),
                                 nn.LeakyReLU()
                                 )
        
        
        self.px3 = nn.Sequential(conv(64, 12, 3, stride=1, padding=1,padding_mode='reflect'),
                                 nn.PixelShuffle(2),
                                 )
        
        #CPR
        self.ecpr = nn.Sequential(nn.Conv2d(3, 16, 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  nn.PixelUnshuffle(2),
                                  nn.PixelUnshuffle(2),
                                  nn.Conv2d(1024, numc, 3, 1, 1,padding_mode='reflect'),
                                  )
        self.dcpr = nn.Sequential(nn.Conv2d(numc, 192, 3, stride=1, padding=1, padding_mode='reflect'), 
                                  nn.PixelShuffle(2),
                                  nn.PixelShuffle(2),
                                  nn.PixelShuffle(2),
                                  )
        
        self.ecpr2 = nn.Sequential(nn.Conv2d(64, 32, 3, 1, 1,padding_mode='reflect'),
                                  nn.PixelUnshuffle(2),
                                  nn.PixelUnshuffle(2),
                                  nn.Conv2d(512, numc, 3, 1, 1,padding_mode='reflect')
                                  )
        self.dcpr2 = nn.Sequential(nn.Conv2d(numc, 192, 3, stride=1, padding=1, padding_mode='reflect'), 
                                  nn.PixelShuffle(2),
                                  nn.PixelShuffle(2),
                                  nn.Conv2d(12, 64, 3, stride=1, padding=1, padding_mode='reflect'), 
                                  )
        
    def encoder(self,x):
        oute = self.enc1(x)     
        x = self.ecpr(x) + self.ecpr2(oute)

        oute = self.enc2(oute)
        oute = self.RinR(oute)
        latent = self.enc3(oute)
        cpr = 0.2*x
        unqlat = self.lat(latent + cpr)

        return unqlat
        
    def decoder(self,unqlat):
        outd = self.px1(unqlat)
        x = self.dcpr2(unqlat)
        unqlat = self.dcpr(unqlat)

        outd = self.RinRdec(outd)
        outd = self.px2(outd) + 0.2*x
        out = self.px3(outd) + 0.2*unqlat 
        out = out.clamp(0,1)
        return out
    
    def forward(self,x):
        unqlat = self.encoder(x)
        out = self.decoder(unqlat)
        return out, unqlat
###############################################################################

def loadCAE(device):
#device = 'cpu'
    autoenc = resConvfmod()
    autoenc = autoenc.to(device)
    autoenc.load_state_dict(torch.load(loaddir, map_location=torch.device(device)))
    autoenc.eval()
    return autoenc

def runCAE(autoenc, valid_dl):
    global counter
    #bpady,bpadx= 0, 0 
    with torch.no_grad():
        for data in valid_dl:
            images, _ = data
            images = images.to(device)
            
            #pad if necessary
            _,_,yy,xx = images.size()
            avgy.append(yy)
            avgx.append(xx)
            if yy%8 !=0 or xx%8!=0:
                imagesp, bpady, bpadx = imgpad(images)
            else:
                imagesp = images
            
            #Compress and decompress
            #'''
            latent = autoenc.encoder(imagesp)
            latnp = tonumpy(latent)
            latint = toint(latnp)
            probs.append(getprobs(latint, int(latnp.max() + abs(float(latnp.min())))))
            latnpby = tobytes(latint)
            #save padding info
            #_,_,yy,xx = images.size()
            if yy%8 !=0 or xx%8!=0:
                latnpby = appendbytes(bpady, bpadx, latnpby)
            #
            savefile(latnpby)
            rec_latnpby = readfile()
            #recover padding info
            if str(rec_latnpby[-3:])[3:4] == 'n':
                rec_latby, pady, padx = popbytes(rec_latnpby)
            else:
                rec_latby = rec_latnpby
            
            rec_latint = frombytes(rec_latby)
            rec_latnp = tofloat(rec_latint)
            rec_latent = fromnumpy(rec_latnp)
            outputV = autoenc.decoder(rec_latent)
            #remove pad
            if str(rec_latnpby[-3:])[3:4] == 'n':
                if pady>0:
                    outputV = outputV[:,:,:-pady,:]
                if padx>0:
                    outputV = outputV[:,:,:,:-padx]
                
            images = images.to('cpu')
            outputV = outputV.to('cpu')
            outputVnp = outputV.detach().numpy()
            outputVnp = np.transpose(outputVnp, (0,2,3,1))
        
            outsave = np.transpose(outputV.detach().numpy(), (0,2,3,1))
            pvalue, svalue, msvalue = psnr(outputV, images).detach(), ssim(outputV,images).detach(), ms_ssim(outputV,images).detach()
            compsize = os.path.getsize(sdirect + 'latby_bz2{0:05d}.bz2'.format(counter))
            print(pvalue,svalue,msvalue, outsave.shape, compsize, counter)
            meanp.append(pvalue)
            means.append(svalue)
            meanms.append(msvalue)
            sizes.append(compsize)

            generated_img = make_grid(outputV)
            save_generator_image(generated_img, sdirect +"img{0:03d}.png".format(counter))
            counter += 1
            #break
    
    npmp = np.array(meanp)
    npms = np.array(means)
    npmms = np.array(meanms)
    npsz =np.array(sizes)
    print(np.mean(npmp), np.mean(npms), np.mean(npmms), np.mean(npsz))
    results.append(np.mean(npmp))
    results.append(np.mean(npms))
    results.append(np.mean(npmms))
    results.append(np.mean(npsz))

plt.figure(1)
print('KODAK')
probs, meanp, means, meanms, sizes = [],[],[],[],[]
avgx, avgy = [], []
autoenc = loadCAE(device)
runCAE(autoenc, valid_dl)

valid_tfms = tt.Compose([tt.Resize((512,768)), tt.ToTensor() ])
valid_ds = ImageFolder(dspath + 'professional_test_2021', valid_tfms)
valid_dl = DataLoader(valid_ds, 1)
counter = 0
print('CLICR')
sdirect = structdir + '/test/CLICR/'   
probs, meanp, means, meanms, sizes = [],[],[],[],[]

plt.figure(2)
runCAE(autoenc, valid_dl)
valid_tfms = tt.Compose([tt.ToTensor() ])
valid_ds = ImageFolder(dspath + 'professional_test_2021', valid_tfms)
valid_dl = DataLoader(valid_ds, 1)

counter = 0
print('CLIC')
sdirect = structdir + '/test/CLIC/'   
probs, meanp, means, meanms, sizes = [],[],[],[],[]
device = 'cpu'

plt.figure(3)
autoenc = loadCAE(device)
avgx, avgy = [], []
runCAE(autoenc, valid_dl)
print(np.mean(np.array(avgx)), np.mean(np.array(avgy)))

#STANFORD
#valid_tfms = tt.Compose([tt.Resize((512,768)), tt.ToTensor() ])
#valid_ds = ImageFolder(dspath + 'cars_tv/v', valid_tfms)
#valid_dl = DataLoader(valid_ds, 1)
#numc = 32
#counter = 0
#print('STF')
#sdirect = structdir + '/test/STANFOUT/'   
probs, meanp, means, meanms, sizes = [],[],[],[],[]

#runCAE()
###############################################################################

print('kdk, clicr, clic')
print(results)   
aaa = np.array(results).reshape((-1,4)).transpose()
df = pd.DataFrame(aaa, columns = ['KODAK','CLICR','CLIC'])
writer = pd.ExcelWriter('E:/demo.xlsx', engine='xlsxwriter')
df.to_excel(writer, sheet_name='Sheet1', index=False)
writer.save()
#'''