├── .gitignore ├── data ├── monarch.bmp ├── monarch_ARCNN.png ├── monarch_FastARCNN.png └── monarch_jpeg_q10.png ├── utils.py ├── model.py ├── example.py ├── README.md ├── dataset.py └── main.py /.gitignore: -------------------------------------------------------------------------------- 1 | .DS_Store 2 | .idea 3 | -------------------------------------------------------------------------------- /data/monarch.bmp: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjn870/ARCNN-pytorch/HEAD/data/monarch.bmp -------------------------------------------------------------------------------- /data/monarch_ARCNN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjn870/ARCNN-pytorch/HEAD/data/monarch_ARCNN.png -------------------------------------------------------------------------------- /data/monarch_FastARCNN.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjn870/ARCNN-pytorch/HEAD/data/monarch_FastARCNN.png -------------------------------------------------------------------------------- /data/monarch_jpeg_q10.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yjn870/ARCNN-pytorch/HEAD/data/monarch_jpeg_q10.png -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | class AverageMeter(object): 2 | def __init__(self): 3 | self.reset() 4 | 5 | def reset(self): 6 | self.val = 0 7 | self.avg = 0 8 | self.sum = 0 9 | self.count = 0 10 | 11 | def update(self, val, n=1): 12 | self.val = val 13 | self.sum += val * n 14 | self.count += n 15 | self.avg = self.sum / self.count 16 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | from torch import nn 2 | 3 | 4 | class ARCNN(nn.Module): 5 | def __init__(self): 6 | super(ARCNN, self).__init__() 7 | self.base = nn.Sequential( 8 | nn.Conv2d(3, 64, kernel_size=9, padding=4), 9 | nn.PReLU(), 10 | nn.Conv2d(64, 32, kernel_size=7, padding=3), 11 | nn.PReLU(), 12 | nn.Conv2d(32, 16, kernel_size=1), 13 | nn.PReLU() 14 | ) 15 | self.last = nn.Conv2d(16, 3, kernel_size=5, padding=2) 16 | 17 | self._initialize_weights() 18 | 19 | def _initialize_weights(self): 20 | for m in self.modules(): 21 | if isinstance(m, nn.Conv2d): 22 | nn.init.normal_(m.weight, std=0.001) 23 | 24 | def forward(self, x): 25 | x = self.base(x) 26 | x = self.last(x) 27 | return x 28 | 29 | 30 | class FastARCNN(nn.Module): 31 | def __init__(self): 32 | super(FastARCNN, self).__init__() 33 | self.base = nn.Sequential( 34 | nn.Conv2d(3, 64, kernel_size=9, stride=2, padding=4), 35 | nn.PReLU(), 36 | nn.Conv2d(64, 32, kernel_size=1), 37 | nn.PReLU(), 38 | nn.Conv2d(32, 32, kernel_size=7, padding=3), 39 | nn.PReLU(), 40 | nn.Conv2d(32, 64, kernel_size=1), 41 | nn.PReLU() 42 | ) 43 | self.last = nn.ConvTranspose2d(64, 3, kernel_size=9, stride=2, padding=4, output_padding=1) 44 | 45 | self._initialize_weights() 46 | 47 | def _initialize_weights(self): 48 | for m in self.modules(): 49 | if isinstance(m, nn.Conv2d): 50 | nn.init.normal_(m.weight, std=0.001) 51 | 52 | def forward(self, x): 53 | x = self.base(x) 54 | x = self.last(x) 55 | return x 56 | -------------------------------------------------------------------------------- /example.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import io 4 | import torch 5 | import torch.backends.cudnn as cudnn 6 | from torchvision import transforms 7 | import PIL.Image as pil_image 8 | from model import ARCNN, FastARCNN 9 | 10 | cudnn.benchmark = True 11 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 12 | 13 | 14 | if __name__ == '__main__': 15 | parser = argparse.ArgumentParser() 16 | parser.add_argument('--arch', type=str, default='ARCNN', help='ARCNN or FastARCNN') 17 | parser.add_argument('--weights_path', type=str, required=True) 18 | parser.add_argument('--image_path', type=str, required=True) 19 | parser.add_argument('--outputs_dir', type=str, required=True) 20 | parser.add_argument('--jpeg_quality', type=int, default=10) 21 | opt = parser.parse_args() 22 | 23 | if not os.path.exists(opt.outputs_dir): 24 | os.makedirs(opt.outputs_dir) 25 | 26 | if opt.arch == 'ARCNN': 27 | model = ARCNN() 28 | elif opt.arch == 'FastARCNN': 29 | model = FastARCNN() 30 | 31 | state_dict = model.state_dict() 32 | for n, p in torch.load(opt.weights_path, map_location=lambda storage, loc: storage).items(): 33 | if n in state_dict.keys(): 34 | state_dict[n].copy_(p) 35 | else: 36 | raise KeyError(n) 37 | 38 | model = model.to(device) 39 | model.eval() 40 | 41 | filename = os.path.basename(opt.image_path).split('.')[0] 42 | 43 | input = pil_image.open(opt.image_path).convert('RGB') 44 | 45 | buffer = io.BytesIO() 46 | input.save(buffer, format='jpeg', quality=opt.jpeg_quality) 47 | input = pil_image.open(buffer) 48 | input.save(os.path.join(opt.outputs_dir, '{}_jpeg_q{}.png'.format(filename, opt.jpeg_quality))) 49 | 50 | input = transforms.ToTensor()(input).unsqueeze(0).to(device) 51 | 52 | with torch.no_grad(): 53 | pred = model(input) 54 | 55 | pred = pred.mul_(255.0).clamp_(0.0, 255.0).squeeze(0).permute(1, 2, 0).byte().cpu().numpy() 56 | output = pil_image.fromarray(pred, mode='RGB') 57 | output.save(os.path.join(opt.outputs_dir, '{}_{}.png'.format(filename, opt.arch))) 58 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AR-CNN, Fast AR-CNN 2 | 3 | This repository is implementation of the "Deep Convolution Networks for Compression Artifacts Reduction".
4 | In contrast with original paper, It use RGB channels instead of luminance channel in YCbCr space and smaller(16) batch size. 5 | 6 | ## Requirements 7 | - PyTorch 8 | - Tensorflow 9 | - tqdm 10 | - Numpy 11 | - Pillow 12 | 13 | **Tensorflow** is required for quickly fetching image in training phase. 14 | 15 | ## Results 16 | 17 | 18 | 19 | 20 | 21 | 22 | 23 | 26 | 29 | 30 | 31 | 32 | 33 | 34 | 35 | 38 | 41 | 42 |
Input
JPEG (Quality 10)
24 |
25 | 		27 |
28 |
AR-CNN
Fast AR-CNN
36 |
37 | 		39 |
40 |
43 | 44 | ## Usages 45 | 46 | ### Train 47 | 48 | When training begins, the model weights will be saved every epoch.
49 | Data augmentation option **--use_augmentation** performs random rescale and rotation.
50 | If you want to train quickly, you should use **--use_fast_loader** option. 51 | 52 | ```bash 53 | python main.py --arch "ARCNN" \ # ARCNN, FastARCNN 54 | --images_dir "" \ 55 | --outputs_dir "" \ 56 | --jpeg_quality 10 \ 57 | --patch_size 24 \ 58 | --batch_size 16 \ 59 | --num_epochs 20 \ 60 | --lr 5e-4 \ 61 | --threads 8 \ 62 | --seed 123 \ 63 | --use_augmentation \ 64 | --use_fast_loader 65 | ``` 66 | 67 | ### Test 68 | 69 | Output results consist of image compressed with JPEG and image with artifacts reduced. 70 | 71 | ```bash 72 | python example --arch "ARCNN" \ # ARCNN, FastARCNN 73 | --weights_path "" \ 74 | --image_path "" \ 75 | --outputs_dir "" \ 76 | --jpeg_quality 10 77 | ``` 78 | -------------------------------------------------------------------------------- /dataset.py: -------------------------------------------------------------------------------- 1 | import os 2 | os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' 3 | 4 | import random 5 | import glob 6 | import io 7 | import numpy as np 8 | import PIL.Image as pil_image 9 | 10 | import tensorflow as tf 11 | config = tf.ConfigProto() 12 | config.gpu_options.allow_growth = True 13 | tf.enable_eager_execution(config=config) 14 | 15 | 16 | class Dataset(object): 17 | def __init__(self, images_dir, patch_size, jpeg_quality, use_augmentation=False, use_fast_loader=False): 18 | self.image_files = sorted(glob.glob(images_dir + '/*')) 19 | self.patch_size = patch_size 20 | self.jpeg_quality = jpeg_quality 21 | self.use_augmentation = use_augmentation 22 | self.use_fast_loader = use_fast_loader 23 | 24 | def __getitem__(self, idx): 25 | if self.use_fast_loader: 26 | label = tf.read_file(self.image_files[idx]) 27 | label = tf.image.decode_jpeg(label, channels=3) 28 | label = pil_image.fromarray(label.numpy()) 29 | else: 30 | label = pil_image.open(self.image_files[idx]).convert('RGB') 31 | 32 | if self.use_augmentation: 33 | # randomly rescale image 34 | if random.random() <= 0.5: 35 | scale = random.choice([0.9, 0.8, 0.7, 0.6]) 36 | label = label.resize((int(label.width * scale), int(label.height * scale)), resample=pil_image.BICUBIC) 37 | 38 | # randomly rotate image 39 | if random.random() <= 0.5: 40 | label = label.rotate(random.choice([90, 180, 270]), expand=True) 41 | 42 | # randomly crop patch from training set 43 | crop_x = random.randint(0, label.width - self.patch_size) 44 | crop_y = random.randint(0, label.height - self.patch_size) 45 | label = label.crop((crop_x, crop_y, crop_x + self.patch_size, crop_y + self.patch_size)) 46 | 47 | # additive jpeg noise 48 | buffer = io.BytesIO() 49 | label.save(buffer, format='jpeg', quality=self.jpeg_quality) 50 | input = pil_image.open(buffer) 51 | 52 | input = np.array(input).astype(np.float32) 53 | label = np.array(label).astype(np.float32) 54 | input = np.transpose(input, axes=[2, 0, 1]) 55 | label = np.transpose(label, axes=[2, 0, 1]) 56 | 57 | # normalization 58 | input /= 255.0 59 | label /= 255.0 60 | 61 | return input, label 62 | 63 | def __len__(self): 64 | return len(self.image_files) 65 | -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import torch 4 | from torch import nn 5 | import torch.optim as optim 6 | import torch.backends.cudnn as cudnn 7 | from torch.utils.data.dataloader import DataLoader 8 | from tqdm import tqdm 9 | from model import ARCNN, FastARCNN 10 | from dataset import Dataset 11 | from utils import AverageMeter 12 | 13 | cudnn.benchmark = True 14 | device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") 15 | 16 | 17 | if __name__ == '__main__': 18 | parser = argparse.ArgumentParser() 19 | parser.add_argument('--arch', type=str, default='ARCNN', help='ARCNN or FastARCNN') 20 | parser.add_argument('--images_dir', type=str, required=True) 21 | parser.add_argument('--outputs_dir', type=str, required=True) 22 | parser.add_argument('--jpeg_quality', type=int, default=10) 23 | parser.add_argument('--patch_size', type=int, default=24) 24 | parser.add_argument('--batch_size', type=int, default=16) 25 | parser.add_argument('--num_epochs', type=int, default=20) 26 | parser.add_argument('--lr', type=float, default=5e-4) 27 | parser.add_argument('--threads', type=int, default=8) 28 | parser.add_argument('--seed', type=int, default=123) 29 | parser.add_argument('--use_augmentation', action='store_true') 30 | parser.add_argument('--use_fast_loader', action='store_true') 31 | opt = parser.parse_args() 32 | 33 | if not os.path.exists(opt.outputs_dir): 34 | os.makedirs(opt.outputs_dir) 35 | 36 | torch.manual_seed(opt.seed) 37 | 38 | if opt.arch == 'ARCNN': 39 | model = ARCNN() 40 | elif opt.arch == 'FastARCNN': 41 | model = FastARCNN() 42 | 43 | model = model.to(device) 44 | criterion = nn.MSELoss() 45 | 46 | optimizer = optim.Adam([ 47 | {'params': model.base.parameters()}, 48 | {'params': model.last.parameters(), 'lr': opt.lr * 0.1}, 49 | ], lr=opt.lr) 50 | 51 | dataset = Dataset(opt.images_dir, opt.patch_size, opt.jpeg_quality, opt.use_augmentation, opt.use_fast_loader) 52 | dataloader = DataLoader(dataset=dataset, 53 | batch_size=opt.batch_size, 54 | shuffle=True, 55 | num_workers=opt.threads, 56 | pin_memory=True, 57 | drop_last=True) 58 | 59 | for epoch in range(opt.num_epochs): 60 | epoch_losses = AverageMeter() 61 | 62 | with tqdm(total=(len(dataset) - len(dataset) % opt.batch_size)) as _tqdm: 63 | _tqdm.set_description('epoch: {}/{}'.format(epoch + 1, opt.num_epochs)) 64 | for data in dataloader: 65 | inputs, labels = data 66 | inputs = inputs.to(device) 67 | labels = labels.to(device) 68 | 69 | preds = model(inputs) 70 | 71 | loss = criterion(preds, labels) 72 | epoch_losses.update(loss.item(), len(inputs)) 73 | 74 | optimizer.zero_grad() 75 | loss.backward() 76 | optimizer.step() 77 | 78 | _tqdm.set_postfix(loss='{:.6f}'.format(epoch_losses.avg)) 79 | _tqdm.update(len(inputs)) 80 | 81 | torch.save(model.state_dict(), os.path.join(opt.outputs_dir, '{}_epoch_{}.pth'.format(opt.arch, epoch))) 82 | --------------------------------------------------------------------------------