Add files via upload

Author: BradleyHWQ

Lossfuncs.py

@@ -0,0 +1,17 @@
+import torch
+
+def mse_loss(input, target):
+    r = input[:,0:1,:,:] - target[:,0:1,:,:]
+    g = input[:,1:2,:,:] - target[:,1:2,:,:]
+    b = input[:,2:3,:,:] - target[:,2:3,:,:]
+    
+    r = torch.mean(r**2)
+    g = torch.mean(g**2)
+    b = torch.mean(b**2)
+    
+    mean = (r + g + b)/3
+   
+    return mean, r,g,b
+
+def parsingLoss(coding, image_size):
+    return torch.sum(torch.abs(coding))/(image_size**2)

Models.py

@@ -0,0 +1,48 @@
+from torch import nn
+
+class autoencoder(nn.Module):
+    def __init__(self):
+        super(autoencoder, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, kernel_size=(5,5))
+        self.maxpool1 = nn.MaxPool2d(kernel_size=(2,2), return_indices=True)
+        self.maxpool2 = nn.MaxPool2d(kernel_size=(2,2), return_indices=True)
+        self.unconv1 = nn.ConvTranspose2d(6,3,kernel_size=(5,5))
+        self.maxunpool1 = nn.MaxUnpool2d(kernel_size=(2,2))
+        self.unmaxunpool2 = nn.MaxUnpool2d(kernel_size=(2,2))
+        
+        self.encoder1 = nn.Sequential(
+            nn.Tanh(),
+            nn.Conv2d(6, 12,kernel_size=(5,5)),
+        )
+        
+        self.encoder2 = nn.Sequential(
+            nn.Tanh(),
+            nn.Conv2d(12, 16, kernel_size=(5,5)),
+            nn.Tanh()
+        )
+        
+        self.decoder2 = nn.Sequential(
+            nn.ConvTranspose2d(16, 12, kernel_size=(5,5)),
+            nn.Tanh()
+        )
+        
+        self.decoder1 = nn.Sequential(
+            nn.ConvTranspose2d(12,6,kernel_size=(5,5)),
+            nn.Tanh(),
+        )
+        
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x,indices1 = self.maxpool1(x)
+        x = self.encoder1(x)
+        x,indices2 = self.maxpool2(x)
+        coding = self.encoder2(x)
+        
+        x = self.decoder2(coding)
+        x = self.unmaxunpool2(x, indices2)
+        x = self.decoder1(x)
+        x = self.maxunpool1(x,indices1)
+        x = self.unconv1(x)
+        output = nn.Tanh()(x)
+        return coding, output

dataloader.py

@@ -0,0 +1,10 @@
+from torch.utils.data import DataLoader
+import torchvision.datasets as dset
+from torchvision import transforms
+
+def DataloaderCompression(dataroot, image_size, batch_size, workers):
+    #dataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([transforms.Resize((image_size, image_size)), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),]))
+    dataset = dset.ImageFolder(root=dataroot, transform=transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),]))
+    
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=workers)
+    return dataloader

decoding.py

@@ -0,0 +1,86 @@
+from matplotlib import pyplot as plt
+import numpy as np
+
+import torch
+import argparse
+from torch.autograd import Variable
+from math import log10
+import torchvision.utils as vutils
+from PIL import Image
+from torchvision import transforms
+#import train_eval
+#from train_eval import to_img
+
+from Models import autoencoder
+from dataloader import DataloaderCompression
+from Lossfuncs import mse_loss, parsingLoss
+
+nb_channls = 3
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--batch_size', type=int, default=8, help='batch size')
+parser.add_argument(
+    '--train', required=True, type=str, help='folder of training images')
+parser.add_argument(
+    '--test', required=True, type=str, help='folder of testing images')
+parser.add_argument(
+    '--max_epochs', type=int, default=50, help='max epochs')
+parser.add_argument('--lr', type=float, default=0.005, help='learning rate')
+# parser.add_argument('--cuda', '-g', action='store_true', help='enables cuda')
+parser.add_argument(
+    '--iterations', type=int, default=100, help='unroll iterations')
+parser.add_argument(
+    '--image_size', type=int, default=150, help='Load image size')
+parser.add_argument('--checkpoint', type=int, default=20, help='save checkpoint after ')
+parser.add_argument('--workers', type=int, default=4, help='unroll iterations')
+parser.add_argument('--weight_decay', type=float, default=0.0005, help='unroll iterations')
+args = parser.parse_args()
+
+def to_img(x):
+    x = 0.5 * (x + 1)
+    x = x.clamp(0, 1)
+    x = x.view(x.size(0), nb_channls, args.image_size, args.image_size)
+    return x
+
+device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
+
+model=torch.load('compressing.pth')
+model.eval()
+
+Dataloader = DataloaderCompression(args.test,args.image_size,args.batch_size,args.workers)
+
+PSNR = []
+Compressing_Ratio = []
+itr = 0
+for data in Dataloader:
+    img, _ = data        
+    img = Variable(img).to(device)
+
+    coding, output = model(img)
+    cyclicloss,r_loss,g_loss,b_loss = mse_loss(output, img)
+
+    PSNR_value = 10*log10(255**2/cyclicloss)
+    PSNR.append(PSNR_value)
+
+    Comp_ratio = coding.size()[1]/img.size()[1]
+    Compressing_Ratio.append(Comp_ratio)
+
+    pic_ = to_img(output.to("cpu").data)
+    #pic = transforms.ToPILImage(pic_)
+            
+    #pic_color = np.transpose(vutils.make_grid(pic.to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0))
+    fig = plt.figure(figsize=(128, 128))
+
+    '''       
+    ax = plt.imshow(np.transpose(vutils.make_grid(pic.to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
+    ax.axes.get_xaxis().set_visible(False)
+    ax.axes.get_yaxis().set_visible(False)
+    '''
+
+    #plt.show(fig)
+    plt.savefig('output/%d.jpg'%itr)
+    itr += 1
+
+print('mean PSNR is %s'%np.mean(PSNR))
+print('mean compression ratio is %s'%np.mean(Compressing_Ratio))

run_test.sh

@@ -0,0 +1,8 @@
+python decoding.py \
+        --batch_size 1 \
+        --train 'Data' \
+        --test 'Data_valid' \
+        --max_epochs 30 \
+        --lr 0.0005 \
+        --iterations 30 \
+        --image_size 128

run_train.sh

@@ -0,0 +1,8 @@
+python train_eval.py \
+        --batch_size 16 \
+        --train 'Data' \
+        --test 'Data_valid' \
+        --max_epochs 30 \
+        --lr 0.0005 \
+        --iterations 30 \
+        --image_size 128

train_eval.py

@@ -0,0 +1,117 @@
+from matplotlib import pyplot as plt
+import numpy as np
+
+import torch
+import torchvision
+from torch import nn
+from torch.autograd import Variable
+from torch.utils.data import DataLoader
+from torchvision import transforms
+import torchvision.datasets as dset
+from torchvision.utils import save_image
+import torchvision.utils as vutils
+from torchsummary import summary
+import argparse
+import sys
+from math import log10
+
+from Models import autoencoder
+from dataloader import DataloaderCompression
+from Lossfuncs import mse_loss, parsingLoss
+
+nb_channls = 3
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--batch_size', type=int, default=8, help='batch size')
+parser.add_argument(
+    '--train', required=True, type=str, help='folder of training images')
+parser.add_argument(
+    '--test', required=True, type=str, help='folder of testing images')
+parser.add_argument(
+    '--max_epochs', type=int, default=50, help='max epochs')
+parser.add_argument('--lr', type=float, default=0.005, help='learning rate')
+# parser.add_argument('--cuda', '-g', action='store_true', help='enables cuda')
+parser.add_argument(
+    '--iterations', type=int, default=100, help='unroll iterations')
+parser.add_argument(
+    '--image_size', type=int, default=150, help='Load image size')
+parser.add_argument('--checkpoint', type=int, default=20, help='save checkpoint after ')
+parser.add_argument('--workers', type=int, default=4, help='unroll iterations')
+parser.add_argument('--weight_decay', type=float, default=0.0005, help='unroll iterations')
+args = parser.parse_args()
+
+device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")
+
+def to_img(x):
+    x = 0.5 * (x + 1)
+    x = x.clamp(0, 1)
+    x = x.view(x.size(0), nb_channls, args.image_size, args.image_size)
+    return x
+
+Dataloader = DataloaderCompression(args.train,args.image_size,args.batch_size,args.workers)
+
+model = autoencoder().to(device)
+criterion = nn.MSELoss()
+
+optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+summary(model, (nb_channls, args.image_size, args.image_size))
+
+# Training Loop. Results will appear every 10th iteration.
+itr = 0
+training_loss = []
+PSNR_list = []
+for epoch in range(args.max_epochs):
+    for data in Dataloader:
+        img, _ = data        
+        img = Variable(img).to(device)
+
+        # Forward
+        coding, output = model(img)
+        cyclicloss,r_loss,g_loss,b_loss = mse_loss(output, img)
+        pLoss = parsingLoss(coding, args.image_size)
+        
+        loss = 5*cyclicloss + 10*pLoss
+
+        PSNR = 10*log10(255**2/cyclicloss)
+
+        # Backprop
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        '''
+        if itr % 10 == 0 and itr < args.iterations:
+            # Log
+            print('iter [{}], whole_loss:{:.4f} cyclic_loss{:.4f} pLoss{:.4f} comp_ratio{:.4f}'
+              .format(itr, loss.data.item(), 5*cyclicloss.data.item(), 10*pLoss.data.item(), PSNR))
+        '''
+        '''
+        if itr % 30 == 0 and itr < args.iterations: 
+            pic = to_img(output.to("cpu").data)
+            
+            fig = plt.figure(figsize=(16, 16))
+            
+            ax = plt.imshow(np.transpose(vutils.make_grid(pic.to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
+            ax.axes.get_xaxis().set_visible(False)
+            ax.axes.get_yaxis().set_visible(False)
+            plt.show(fig)
+            
+            
+            compress_ratio.append(comp_ratio)
+        '''
+        training_loss.append(loss)
+        PSNR_list.append(PSNR)
+        itr += 1
+    
+        print('epoch [{}/{}], loss:{:.4f}, cyclic_loss{:.4f} pLoss{:.4f} PSNR{:.4f}'
+            .format(epoch + 1, args.max_epochs, loss.data.item(), 5*cyclicloss.data.item(), 10*pLoss.data.item(), PSNR))
+
+    if epoch % 10 == 0:
+        torch.save(model, 'Compressing_{%d}.pth'%epoch)
+
+plt.plot(training_loss, label='Training loss')
+plt.plot(PSNR, label='PSNR')
+plt.legend(frameon=False)
+plt.savefig("Train.png")
+plt.show()