diff --git a/README.md b/README.md index 706d398..61839b1 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ We have a list of candidate papers to implement: https://github.com/chainer/mode - Neural Relational Inference for Interacting Systems [[paper](https://arxiv.org/abs/1802.04687)] [[code](https://github.com/chainer/models/tree/master/nri)] - SiamRPN and SiamMask [[paper](https://arxiv.org/abs/1812.05050)] [[code](https://github.com/STVIR/pysot)] - Learning to learn by gradient descent by gradient descent [[paper](https://arxiv.org/abs/1606.04474)] [[code](https://github.com/chainer/models/tree/master/learning_to_learn)] +- Learning to Simplify: Fully Convolutional Networks for Rough Sketch Cleanup [[paper](http://www.f.waseda.jp/hfs/SimoSerraSIGGRAPH2016.pdf)] [[code](https://github.com/chainer/models/tree/master/simplifying_rough_sketches)] ## License MIT License (see `LICENSE` file). diff --git a/simplifying_rough_sketches/README.md b/simplifying_rough_sketches/README.md new file mode 100644 index 0000000..ad89130 --- /dev/null +++ b/simplifying_rough_sketches/README.md @@ -0,0 +1,133 @@ +# Rough Sketch Simplification using FCNN in PyTorch + +This repository contains code of the paper [Learning to Simplify: Fully Convolutional Networks for Rough Sketch Cleanup](http://www.f.waseda.jp/hfs/SimoSerraSIGGRAPH2016.pdf) which is tested and trained on custom datasets. It is based on Chainer. + +## Overview + +The paper presents novel technique to simplify sketch drawings based on learning a series of convolution operators. Image of any dimension can be fed into the network, and it outputs the image of same dimension as the input image. + +![model](images/model.png) + +The architecture consists of encoder and a decoder, the first part acts as an encoder and spatially compresses the image, the second part, processes and extracts the essential lines from the image, and the third and last part acts as a decoder which converts the small more simple representation to an grayscale image of the same resolution as the input. This is all done using convolutions. +The down- and up-convolution architecture may seem similar to a simple filter banks. However, it is important to realize that the number of channels is much larger where resolution is lower, e.g., 1024 where the size is 1/8. This ensures that information that leads to clean lines is carried through the low-resolution part; the network is trained to choose which information to carry by the encoder- decoder architecture. Padding is used to compensate for the kernel size and ensure the output is the same size as the input when a stride of 1 is used. Pooling layers are replaced by convolutional layers with increased strides to lower the resolution from the previous layer. + + + +## Contents +- [Rough Sketch Simplification using FCNN in PyTorch](#rough-sketch-simplification-using-fcnn-in-pytorch) + - [Overview](#overview) + - [Contents](#contents) + - [1. Setup Instructions and Dependencies](#1-setup-instructions-and-dependencies) + - [2. Dataset](#2-dataset) + - [3. Training the model](#3-training-the-model) + - [5. Model Architecture](#5-model-architecture) + - [6. Observations](#6-observations) + - [Training](#training) + - [Predicitons](#predicitons) + - [Loss](#loss) + - [7. Repository overview](#7-repository-overview) + + +## 1. Setup Instructions and Dependencies + +Clone the repositiory on your local machine. + + +Start a virtual environment using python3 +``` Batchfile +virtualenv env +``` + + +Install the dependencies +``` Batchfile +pip install -r requirements.txt +``` + +You can also use google collab notebook. + + +## 2. Dataset + +The authors have not provided dataset for the paper. So I created my own. I have uploaded the dataset on drive, the link to which can be found [here](https://drive.google.com/open?id=14NQTqITAiw8o-JgdnumQ-K0asLRwJy7q). Feel free to use it. + +Create two folders inside the root directory of dataset, `Input` and `Taget` and place the images inside the corresponding directory. It is important to keep the names same for both input and target images. + +## 3. Training the model + +To train the model, run + +```Batchfile +python main.py --train=True +``` + +optional arguments: + + | argument | default | desciption| + | --- | --- | --- | + | -h, --help | None | show help message and exit | + | --gpu_id GPU_ID, -g GPU_ID | -1 | GPU ID (negative value indicates CPU) | + | --out OUT, -o OUT |result | Directory to output the result | + | --batch_size BATCH_SIZE, -b BATCH_SIZE | 8 | Batch Size | + | --height HEIGHT, -ht HEIGHT | 64 | height of the image to resize to | + | --width WIDTH, -wd WIDTH | 64 | width of the image to resize to | + | --samples SAMPLES | False | See sample training images | + | --num_epochs NUM_EPOCHS | 75 | Number of epochs to train on | + | --train TRAIN | True | train the model | + | --root ROOT | . | Root Directory for Input and Target images. | + | --n_folds N_FOLDS | 7 | Number of folds in k-fold cross validation. | + | --save_model SAVE_MODEL | True | Save model after training. | + | --load_model LOAD_MODEL | None | Path to existing model. | + | --predict PREDICT | None | Path of rough sketch to simplify using existing model | + +## 5. Model Architecture + +![archi](images/archi.png) + +## 6. Observations + + +### Training + +| Epoch | Prediction | +| --- | --- | +| 2 | ![epoch2](pred/2.png) | +| 60 | ![epoch40](pred/60.png) | +| 100 | ![epoch80](pred/100.png) | +| 140 | ![epoch120](pred/140.png) | + +### Predicitons + +![pred3](pred/pred3.png) +![pred2](pred/pred8.png) +![pred1](pred/pred1.png) + + +### Loss + +![loss](images/loss.png) + +## 7. Repository overview + +This repository contains the following files and folders + +1. **images**: Contains resourse images. + +2. **pred**: Contains prediction images. + +3. `dataset.py`: code for dataset generation. + +4. `model.py`: code for model as described in the paper. + +5. `predict.py`: function to simplify image using model. + +6. `read_data.py`: code to read images. + +7. `utils.py`: Contains helper functions. + +8. `train_val.py`: function to train and validate models. + +9. `main.py`: contains main code to run the model. + +10. `requirements.txt`: Lists dependencies for easy setup in virtual environments. + diff --git a/simplifying_rough_sketches/dataset.py b/simplifying_rough_sketches/dataset.py new file mode 100644 index 0000000..931285a --- /dev/null +++ b/simplifying_rough_sketches/dataset.py @@ -0,0 +1,35 @@ +from PIL import Image +import numpy as np + +from chainer.dataset import dataset_mixin + +class CustomDataset(dataset_mixin.DatasetMixin): + + def __init__(self, input_path, target_path, height=424, width=424): + self.input_img = input_path + self.target_img = target_path + self.height = height + self.width = width + + def __len__(self): + return len(self.input_img) + + def preprocess(self, img_path): + img = Image.open(img_path) + img = img.convert('L') + img = img.resize((self.height, self.width)) + img = np.asarray(img)/255.0 + img = img.reshape(1,self.height, self.width) + return img + + def __getitem__(self, index): + X = self.preprocess(self.input_img[index]) + y = self.preprocess(self.target_img[index]) + + return X, y + + def get_example(self, index): + X = self.preprocess(self.input_img[index]) + y = self.preprocess(self.target_img[index]) + + return X, y diff --git a/simplifying_rough_sketches/images/archi.png b/simplifying_rough_sketches/images/archi.png new file mode 100644 index 0000000..cfad1d8 Binary files /dev/null and b/simplifying_rough_sketches/images/archi.png differ diff --git a/simplifying_rough_sketches/images/loss.png b/simplifying_rough_sketches/images/loss.png new file mode 100644 index 0000000..703b42c Binary files /dev/null and b/simplifying_rough_sketches/images/loss.png differ diff --git a/simplifying_rough_sketches/images/model.png b/simplifying_rough_sketches/images/model.png new file mode 100644 index 0000000..e7154f5 Binary files /dev/null and b/simplifying_rough_sketches/images/model.png differ diff --git a/simplifying_rough_sketches/main.py b/simplifying_rough_sketches/main.py new file mode 100644 index 0000000..9768ac5 --- /dev/null +++ b/simplifying_rough_sketches/main.py @@ -0,0 +1,151 @@ +import chainer + +import time +import argparse +import os + +from PIL import Image +import numpy as np +import matplotlib.pyplot as plt + +from model import Net +from dataset import CustomDataset +from read_data import get_data +from utils import samples +from train_val import train, validate +from utils import argument_parser, create_directory +from predict import predict + +def main(): + + args = argument_parser() + + gpu_id = args.gpu_id + # root = args.root + height = args.height + width = args.width + BATCHSIZE = args.batch_size + N_FOLD = args.n_folds + num_epochs = args.num_epochs + + out = args.out + create_directory(out) + + root = args.root + Input = os.path.join(root, 'Input') + Target = os.path.join(root,'Target') + # height = width = 424 + + input_images, target_images = get_data(Input, Target) + + + if args.samples: + show_samples(input_images) + show_samples(target_images) + + + if args.predict and args.load_model: + + model_path = args.load_model + img_path = args.predict + + model = Net() + + from chainer import serializers + serializers.load_npz(model_path, model) + print ('Model Loaded') + + + img = predict(model, img_path, height, width, gpu_id) + pilimg = Image.fromarray(np.uint8(img)) + pilimg.save(os.path.join(out,f'predict{hash(img_path)}.png')) + print (f'Image saved in {out} directory.') + + + elif args.train: + + print ('Training Started.') + + model = Net() + + if gpu_id >= 0: + chainer.backends.cuda.get_device_from_id(gpu_id).use() + model.to_gpu() + print ('Switched to CUDA backend') + else: + print ('Training on CPU') + + + optimizer = chainer.optimizers.AdaDelta(rho=0.9) + optimizer.setup(model) + optimizer.use_cleargrads(True) + + dataset = CustomDataset(input_path=input_images, target_path=target_images, height=height, width=width) + train_val_set = chainer.datasets.get_cross_validation_datasets_random(dataset=dataset, n_fold=N_FOLD, seed=None) + + total_loss_val, total_loss_train = [],[] + + import time + since = time.time() + epoch_num = num_epochs + best_val_loss = 1000 + + for fold, train_val in enumerate(train_val_set): + + train_set, val_set = train_val + + train_loader = chainer.iterators.MultithreadIterator(dataset=train_set, batch_size=BATCHSIZE, shuffle=True, repeat=False) + val_loader = chainer.iterators.MultithreadIterator(dataset=val_set, batch_size=BATCHSIZE, shuffle=True, repeat=False) + + for epoch in range(1, epoch_num+1): + # chainer.cuda.memory_pool.free_all_blocks() + avg_loss_train, loss_train = train(train_loader, model, optimizer, epoch, gpu_id) + loss_val = validate(val_loader, model, optimizer, epoch, gpu_id) + total_loss_val.append(loss_val) + total_loss_train.append(avg_loss_train) + + if loss_val < best_val_loss: + best_val_loss = loss_val + print('*****************************************************') + print(f'best record: [Fold {fold}] [epoch {epoch}], [val loss {loss_val:.5f}]') + print('*****************************************************') + + if epoch%1 == 0: + print ('Fold: {0} Epoch: {1}'.format(fold,epoch)) + samples(input_images[0], target_images[0], model, out, height, width, gpu_id) + img = predict(model, input_images[0], height, width, gpu_id) + pilimg = Image.fromarray(np.uint8(img)) + pilimg.save(os.path.join(out,str(fold)+'_'+str(epoch)+'.png')) + + end = time.time() + + print ('Training Completed') + if args.save_model: + print('Saving model in models directory') + from chainer import serializers + create_directory('models') + serializers.save_npz(os.path.join('models','simplifying_rough_sketches.model'), model) + + print ('Time Taken: ',end-since) + fig = plt.figure(num = 2) + fig1 = fig.add_subplot() + fig1.plot(total_loss_train, label = 'training loss') + fig1.plot(total_loss_val, label = 'validation loss') + plt.legend(loc='upper left') + plt.savefig(os.path.join(out,'loss.png')) + plt.close(fig) + + + try: + for i in range (5): + import random + k = random.randint(1,63) + samples(input_images[k], target_images[k], model, out) + except Exception as e: + pass + + else: + print ('Exiting') + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/simplifying_rough_sketches/model.py b/simplifying_rough_sketches/model.py new file mode 100644 index 0000000..b920143 --- /dev/null +++ b/simplifying_rough_sketches/model.py @@ -0,0 +1,222 @@ +import chainer +import chainer.functions as F +import chainer.links as L + + +class Net(chainer.Chain): + + def __init__(self): + + super(Net, self).__init__() + + with self.init_scope(): + + self.down_conv1 = chainer.Sequential( + L.Convolution2D(in_channels=1, + out_channels=48, + ksize=5, + stride=2, + pad=2), + L.BatchNormalization(48), + F.relu, + + L.Convolution2D(in_channels=48, + out_channels=128, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(128), + F.relu, + + L.Convolution2D(in_channels=128, + out_channels=128, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(128), + F.relu, + ) + + + self.down_conv2 = chainer.Sequential( + L.Convolution2D(in_channels=128, + out_channels=256, + ksize=3, + stride=2, + pad=1), + L.BatchNormalization(256), + F.relu, + + L.Convolution2D(in_channels=256, + out_channels=256, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(256), + F.relu, + + L.Convolution2D(in_channels=256, + out_channels=256, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(256), + F.relu + ) + + + self.flat_conv = chainer.Sequential( + L.Convolution2D(in_channels=256, + out_channels=256, + ksize=3, + stride=2, + pad=1), + L.BatchNormalization(256), + F.relu, + + L.Convolution2D(in_channels=256, + out_channels=512, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(512), + F.relu, + + L.Convolution2D(in_channels=512, + out_channels=1024, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(1024), + F.relu, + + L.Convolution2D(in_channels=1024, + out_channels=1024, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(1024), + F.relu, + + L.Convolution2D(in_channels=1024, + out_channels=1024, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(1024), + F.relu, + + L.Convolution2D(in_channels=1024, + out_channels=1024, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(1024), + F.relu, + + L.Convolution2D(in_channels=1024, + out_channels=512, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(512), + F.relu, + + L.Convolution2D(in_channels=512, + out_channels=256, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(256), + F.relu + ) + + + self.up_conv1 = chainer.Sequential( + L.Deconvolution2D(in_channels=256, + out_channels=256, + ksize=4, + stride=2, + pad=1), + L.BatchNormalization(256), + F.relu, + + L.Convolution2D(in_channels=256, + out_channels=256, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(256), + F.relu, + + L.Convolution2D(in_channels=256, + out_channels=128, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(128), + F.relu, + ) + + + self.up_conv2 = chainer.Sequential( + L.Deconvolution2D(in_channels=128, + out_channels=128, + ksize=4, + stride=2, + pad=1), + L.BatchNormalization(128), + F.relu, + + L.Convolution2D(in_channels=128, + out_channels=128, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(128), + F.relu, + + L.Convolution2D(in_channels=128, + out_channels=48, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(48), + F.relu, + ) + + + self.up_conv3 = chainer.Sequential( + L.Deconvolution2D(in_channels=48, + out_channels=48, + ksize=4, + stride=2, + pad=1), + L.BatchNormalization(48), + F.relu, + + L.Convolution2D(in_channels=48, + out_channels=24, + ksize=3, + stride=1, + pad=1), + L.BatchNormalization(24), + F.relu, + + L.Convolution2D(in_channels=24, + out_channels=1, + ksize=3, + stride=1, + pad=1), + F.sigmoid + ) + + + def forward(self, x): + x = self.down_conv1(x) + x = self.down_conv2(x) + x = self.flat_conv(x) + x = self.up_conv1(x) + x = self.up_conv2(x) + out = self.up_conv3(x) + return out diff --git a/simplifying_rough_sketches/pred/100.png b/simplifying_rough_sketches/pred/100.png new file mode 100644 index 0000000..d6b0ef6 Binary files /dev/null and b/simplifying_rough_sketches/pred/100.png differ diff --git a/simplifying_rough_sketches/pred/120.png b/simplifying_rough_sketches/pred/120.png new file mode 100644 index 0000000..adf9aa4 Binary files /dev/null and b/simplifying_rough_sketches/pred/120.png differ diff --git a/simplifying_rough_sketches/pred/140.png b/simplifying_rough_sketches/pred/140.png new file mode 100644 index 0000000..8113198 Binary files /dev/null and b/simplifying_rough_sketches/pred/140.png differ diff --git a/simplifying_rough_sketches/pred/2.png b/simplifying_rough_sketches/pred/2.png new file mode 100644 index 0000000..8f60c77 Binary files /dev/null and b/simplifying_rough_sketches/pred/2.png differ diff --git a/simplifying_rough_sketches/pred/20.png b/simplifying_rough_sketches/pred/20.png new file mode 100644 index 0000000..d42f4f2 Binary files /dev/null and b/simplifying_rough_sketches/pred/20.png differ diff --git a/simplifying_rough_sketches/pred/40.png b/simplifying_rough_sketches/pred/40.png new file mode 100644 index 0000000..1857e6f Binary files /dev/null and b/simplifying_rough_sketches/pred/40.png differ diff --git a/simplifying_rough_sketches/pred/60.png b/simplifying_rough_sketches/pred/60.png new file mode 100644 index 0000000..d6be742 Binary files /dev/null and b/simplifying_rough_sketches/pred/60.png differ diff --git a/simplifying_rough_sketches/pred/80.png b/simplifying_rough_sketches/pred/80.png new file mode 100644 index 0000000..6f374af Binary files /dev/null and b/simplifying_rough_sketches/pred/80.png differ diff --git a/simplifying_rough_sketches/pred/pred1.png b/simplifying_rough_sketches/pred/pred1.png new file mode 100644 index 0000000..dcf14db Binary files /dev/null and b/simplifying_rough_sketches/pred/pred1.png differ diff --git a/simplifying_rough_sketches/pred/pred2.png b/simplifying_rough_sketches/pred/pred2.png new file mode 100644 index 0000000..7ecd638 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred2.png differ diff --git a/simplifying_rough_sketches/pred/pred3.png b/simplifying_rough_sketches/pred/pred3.png new file mode 100644 index 0000000..28d4d48 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred3.png differ diff --git a/simplifying_rough_sketches/pred/pred4.png b/simplifying_rough_sketches/pred/pred4.png new file mode 100644 index 0000000..992e0d2 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred4.png differ diff --git a/simplifying_rough_sketches/pred/pred5.png b/simplifying_rough_sketches/pred/pred5.png new file mode 100644 index 0000000..cbec1e1 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred5.png differ diff --git a/simplifying_rough_sketches/pred/pred6.png b/simplifying_rough_sketches/pred/pred6.png new file mode 100644 index 0000000..c7bfe65 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred6.png differ diff --git a/simplifying_rough_sketches/pred/pred7.png b/simplifying_rough_sketches/pred/pred7.png new file mode 100644 index 0000000..e1f4e2f Binary files /dev/null and b/simplifying_rough_sketches/pred/pred7.png differ diff --git a/simplifying_rough_sketches/pred/pred8.png b/simplifying_rough_sketches/pred/pred8.png new file mode 100644 index 0000000..84d18b1 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred8.png differ diff --git a/simplifying_rough_sketches/pred/pred9.png b/simplifying_rough_sketches/pred/pred9.png new file mode 100644 index 0000000..6702be7 Binary files /dev/null and b/simplifying_rough_sketches/pred/pred9.png differ diff --git a/simplifying_rough_sketches/pred/prred10.png b/simplifying_rough_sketches/pred/prred10.png new file mode 100644 index 0000000..c00d59a Binary files /dev/null and b/simplifying_rough_sketches/pred/prred10.png differ diff --git a/simplifying_rough_sketches/predict.py b/simplifying_rough_sketches/predict.py new file mode 100644 index 0000000..0ca1ea3 --- /dev/null +++ b/simplifying_rough_sketches/predict.py @@ -0,0 +1,30 @@ +import chainer +from chainer import Variable + +import numpy as np +import cv2 + +from utils import preprocess + +def predict(model, img_path, height, width, gpu_id): + with chainer.no_backprop_mode(): + in_shape = np.asarray(cv2.imread(img_path)).shape + img = preprocess(img_path, height, width) + fin_shape = np.asarray(img).shape + img = img.reshape(1, 1, fin_shape[0],fin_shape[1]) + + img = Variable(img.astype('float32')) + + if gpu_id >= 0: + chainer.backends.cuda.get_device_from_id(gpu_id).use() + img.to_gpu() + + output = model(img) + output.to_cpu() + img = np.array(output.data) * 255 + # img = img.astype('uint8') + img = img.squeeze().squeeze() + # print (img.shape) + + img = cv2.resize(img, (in_shape[1],in_shape[0]),interpolation = cv2.INTER_AREA) + return img diff --git a/simplifying_rough_sketches/read_data.py b/simplifying_rough_sketches/read_data.py new file mode 100644 index 0000000..34d7261 --- /dev/null +++ b/simplifying_rough_sketches/read_data.py @@ -0,0 +1,14 @@ +import os + +def get_data(Input, Target): + Targets =sorted(os.listdir(Target)) + Inputs = sorted(os.listdir(Input)) + input_images, target_images = [], [] + for input_image, target_image in zip(Inputs, Targets): + image_path = os.path.join(Input, input_image) + target_path = os.path.join(Target, target_image) + + input_images.append(image_path) + target_images.append(target_path) + + return input_images, target_images diff --git a/simplifying_rough_sketches/requirements.txt b/simplifying_rough_sketches/requirements.txt new file mode 100644 index 0000000..d8b6425 --- /dev/null +++ b/simplifying_rough_sketches/requirements.txt @@ -0,0 +1,6 @@ +'cupy-cuda100>=6.4.0,<7.0.0' +chainer==6.5.0 +numpy==1.17.3 +matplotlib==3.1.1 +opencv-python==3.4.7.28 +Pillow==4.3.0 diff --git a/simplifying_rough_sketches/train_val.py b/simplifying_rough_sketches/train_val.py new file mode 100644 index 0000000..c77641a --- /dev/null +++ b/simplifying_rough_sketches/train_val.py @@ -0,0 +1,70 @@ +from utils import AverageMeter, length + +import chainer +from chainer import Variable +import chainer.functions as F + + +def train(train_loader, model, optimizer, epoch, gpu_id): + train_loader.reset() + + train_loss = AverageMeter() + + curr_iter = (epoch - 1) * length(train_loader) + loss_train = [] + + + for i, data in enumerate(train_loader): + + model.cleargrads() + + batch_img, targets = chainer.dataset.concat_examples(data) + + inputs = Variable(batch_img.astype('float32')) + targets = Variable(targets.astype('float32')) + + if gpu_id >= 0: + chainer.backends.cuda.get_device_from_id(gpu_id).use() + inputs.to_gpu() + targets.to_gpu() + + outputs = model(inputs) + loss = F.mean_squared_error(outputs, targets) + train_loss.update(loss.item()) + loss.backward() + optimizer.update() + curr_iter += 1 + if epoch % 1 == 0: + if (i + 1) % 8 == 0: + print('[epoch %d], [iter %d / %d], [train loss %.5f]' % ( + epoch, i + 1, length(train_loader), train_loss.avg)) + loss_train.append(train_loss.avg) + + return train_loss.avg, loss_train + + +def validate(val_loader, model, optimizer, epoch, gpu_id): + val_loss = AverageMeter() + + with chainer.no_backprop_mode(): + val_loader.reset() + for _, data in enumerate(val_loader): + images, targets = chainer.dataset.concat_examples(data) + + images = Variable(images.astype('float32')) + targets = Variable(targets.astype('float32')) + + if gpu_id >= 0: + chainer.backends.cuda.get_device_from_id(gpu_id).use() + images.to_gpu() + targets.to_gpu() + + outputs = model(images) + loss = F.mean_squared_error(outputs, targets) + val_loss.update(loss.item()) + + print('------------------------------------------------------------') + print('[epoch %d], [val loss %.5f]' % (epoch, val_loss.avg)) + print('------------------------------------------------------------') + return val_loss.avg + diff --git a/simplifying_rough_sketches/utils.py b/simplifying_rough_sketches/utils.py new file mode 100644 index 0000000..e25beed --- /dev/null +++ b/simplifying_rough_sketches/utils.py @@ -0,0 +1,84 @@ +from PIL import Image +import matplotlib.pyplot as plt +import numpy as np + +import math +import argparse +import os +import time + +def length(train_loader): return math.ceil(len(train_loader.dataset)/train_loader.batch_size) + + +# this function is used during training process, to calculation the loss and accuracy +class AverageMeter(object): + def __init__(self): + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + +def preprocess(img_path, height, width): + img = Image.open(img_path) + img = img.convert('L') + img = img.resize((height, width)) + img = np.asarray(img) + return img + + +def samples(input_image, target_image, model, out, height, width, gpu_id): + from predict import predict + + _,figs = plt.subplots(1,3) + plt.subplots_adjust(wspace=0.2, hspace=0.2) + plt.rcParams['figure.figsize'] = [15, 15] + figs[0].set_title('Input') + figs[0].imshow(Image.open(input_image),cmap = 'gray') + + figs[1].set_title('Prediction') + figs[1].imshow(predict(model, input_image, height, width, gpu_id),cmap='gray') + + figs[2].set_title('Target') + figs[2].imshow(Image.open(target_image),cmap = 'gray') + + plt.savefig(os.path.join(out,f'sample{time.time()}.png')) + + + +def argument_parser(): + parser = argparse.ArgumentParser(description='Rough Sketch Simplification') + parser.add_argument('--gpu_id', '-g', type=int, default=-1, help='GPU ID (negative value indicates CPU)') + parser.add_argument('--out', '-o', type=str, default='result', help='Directory to output the result') + parser.add_argument('--batch_size', '-b', type=int, default=8, help='Batch Size') + parser.add_argument('--height', '-ht', type=int, default=64, help='height of the image to resize to') + parser.add_argument('--width', '-wd', type=int, default=64, help='width of the image to resize to') + parser.add_argument('--samples', type=bool, default=False, help='See sample training images') + parser.add_argument('--num_epochs', type=int, default=75, help='Number of epochs to train on') + parser.add_argument('--train', default=True, type=bool, help='train the model') + parser.add_argument('--root', default='.', type=str, help='Root Directory for Input and Target images.') + parser.add_argument('--n_folds', default=7, type=int, help='Number of folds in k-fold cross validation.') + parser.add_argument('--save_model', default=True, type=bool, help='Save model after training.') + parser.add_argument('--load_model', default=None, type=str, help='Load existing model.') + parser.add_argument('--predict', default=None, type=str, help='Path of rough sketch to simplify using existing model') + + args = parser.parse_args() + + return args + +def create_directory(path): + try: + # Create target Directory + os.mkdir(path) + print("Directory " , path , " Created ") + except FileExistsError: + print("Directory " , path , " already exists")