Skip to content

Commit

Permalink
format
Browse files Browse the repository at this point in the history
  • Loading branch information
@wuhuikai
wuhuikai committed Sep 24, 2019
1 parent 0047f43 commit e166651
Show file tree
Hide file tree
Showing 11 changed files with 196 additions and 171 deletions.
9 changes: 5 additions & 4 deletions crop_aligned_images.py
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import os
import glob
import argparse
import glob
import os

from skimage.io import imread, imsave

Expand All @@ -9,7 +9,8 @@ def main():
parser = argparse.ArgumentParser(description='Cropping aligned images')
parser.add_argument('--data_root', help='Path for aligned images')
parser.add_argument('--result_folder', default='cropped_images', help='Path for cropped images')
parser.add_argument('--bbox_path', default='DataBase/TransientAttributes/bbox.txt', help='Path for bounding-box txt')
parser.add_argument('--bbox_path', default='DataBase/TransientAttributes/bbox.txt',
help='Path for bounding-box txt')
args = parser.parse_args()

# Init mask folder
Expand All @@ -30,7 +31,7 @@ def main():
for data_root in images:
mask = imread(data_root)
cropped_mask = mask[sx:ex, sy:ey]

mask_name = os.path.basename(data_root)
imsave(os.path.join(args.result_folder, name, mask_name), cropped_mask)

Expand Down
41 changes: 22 additions & 19 deletions dataset.py
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,12 @@
import os
import glob
import math
import numpy
import os

import numpy
from chainer.dataset import dataset_mixin
from skimage.io import imread
from skimage.transform import resize

from chainer.dataset import dataset_mixin


class H5pyDataset(dataset_mixin.DatasetMixin):
def __init__(self, path, which_set='train', load_size=None, crop_size=None, dtype=numpy.float32):
Expand All @@ -23,19 +22,20 @@ def __len__(self):

def get_example(self, i):
handle = self._data_set.open()
data = self._data_set.get_data(handle, slice(i, i+1))
data = self._data_set.get_data(handle, slice(i, i + 1))
self._data_set.close(handle)

im = numpy.squeeze(data[0])

w, h, _ = im.shape
min_size = min(w, h)
ratio = self._load_size/min_size
rw, rh = int(math.ceil(w*ratio)), int(math.ceil(h*ratio))
ratio = self._load_size / min_size
rw, rh = int(math.ceil(w * ratio)), int(math.ceil(h * ratio))
im = resize(im, (rw, rh), order=1, mode='constant')

sx, sy = numpy.random.random_integers(0, rw-self._crop_size), numpy.random.random_integers(0, rh-self._crop_size)
im = im[sx:sx+self._crop_size, sy:sy+self._crop_size,:]*2 - 1
sx, sy = numpy.random.random_integers(0, rw - self._crop_size), numpy.random.random_integers(0,
rh - self._crop_size)
im = im[sx:sx + self._crop_size, sy:sy + self._crop_size, :] * 2 - 1

im = numpy.asarray(numpy.transpose(im, (2, 0, 1)), dtype=self._dtype)

Expand All @@ -44,14 +44,14 @@ def get_example(self, i):

class BlendingDataset(dataset_mixin.DatasetMixin):
def __init__(self, total_examples, folders, root, ratio, load_size, crop_size, dtype=numpy.float32):
imgs_per_folder = {folder:glob.glob(os.path.join(root, folder, '*')) for folder in folders}
imgs_per_folder = {folder: glob.glob(os.path.join(root, folder, '*')) for folder in folders}
self._len = total_examples

self._dtype = dtype
self._load_size = load_size
self._crop_size = crop_size
self._size = int(self._crop_size*ratio)
self._sx = self._crop_size//2 - self._size//2
self._size = int(self._crop_size * ratio)
self._sx = self._crop_size // 2 - self._size // 2

self._imgs = []
for _ in range(self._len):
Expand All @@ -64,26 +64,29 @@ def __len__(self):

def _crop(self, im, rw, rh, sx, sy):
im = resize(im, (rw, rh), order=1, preserve_range=False, mode='constant')
im = im[sx:sx+self._crop_size, sy:sy+self._crop_size,:]*2 - 1
im = im[sx:sx + self._crop_size, sy:sy + self._crop_size, :] * 2 - 1
im = numpy.transpose(im, (2, 0, 1)).astype(self._dtype)

return im

def get_example(self, i):
obj_path, bg_path = self._imgs[i]
obj = imread(obj_path)
bg = imread(bg_path)
bg = imread(bg_path)

w, h, _ = obj.shape
min_size = min(w, h)
ratio = self._load_size/min_size
rw, rh = int(math.ceil(w*ratio)), int(math.ceil(h*ratio))
sx, sy = numpy.random.random_integers(0, rw-self._crop_size), numpy.random.random_integers(0, rh-self._crop_size)
ratio = self._load_size / min_size
rw, rh = int(math.ceil(w * ratio)), int(math.ceil(h * ratio))
sx, sy = numpy.random.random_integers(0, rw - self._crop_size), numpy.random.random_integers(0,
rh - self._crop_size)

obj_croped = self._crop(obj, rw, rh, sx, sy)
bg_croped = self._crop(bg, rw, rh, sx, sy)
bg_croped = self._crop(bg, rw, rh, sx, sy)

copy_paste = bg_croped.copy()
copy_paste[:, self._sx:self._sx+self._size, self._sx:self._sx+self._size] = obj_croped[:, self._sx:self._sx+self._size, self._sx:self._sx+self._size]
copy_paste[:, self._sx:self._sx + self._size, self._sx:self._sx + self._size] = obj_croped[:,
self._sx:self._sx + self._size,
self._sx:self._sx + self._size]

return copy_paste, bg_croped
89 changes: 47 additions & 42 deletions gp_gan.py
View file
Original file line number Diff line number Diff line change
@@ -1,35 +1,34 @@
import math

import numpy as np

from skimage.filters import gaussian, sobel_h, sobel_v, scharr_h, scharr_v, roberts_pos_diag, roberts_neg_diag, prewitt_h, prewitt_v
from skimage.transform import resize

from scipy.optimize import minimize
from scipy.ndimage import correlate
from scipy.fftpack import dct, idct

import chainer
import chainer.functions as F
import numpy as np
from chainer import cuda, Variable

from scipy.fftpack import dct, idct
from scipy.ndimage import correlate
from scipy.optimize import minimize
from skimage.filters import gaussian, sobel_h, sobel_v, scharr_h, scharr_v, roberts_pos_diag, roberts_neg_diag, \
prewitt_h, prewitt_v
from skimage.transform import resize

################## Gradient Operator #########################
normal_h = lambda im: correlate(im, np.asarray([[0, -1, 1]]), mode='nearest')
normal_h = lambda im: correlate(im, np.asarray([[0, -1, 1]]), mode='nearest')
normal_v = lambda im: correlate(im, np.asarray([[0, -1, 1]]).T, mode='nearest')

gradient_operator = {
'normal': (normal_h, normal_v),
'sobel': (sobel_h, sobel_v),
'scharr': (scharr_h, scharr_v),
'normal': (normal_h, normal_v),
'sobel': (sobel_h, sobel_v),
'scharr': (scharr_h, scharr_v),
'roberts': (roberts_pos_diag, roberts_neg_diag),
'prewitt': (prewitt_h, prewitt_v)
}


###########################################################


def preprocess(im):
im = np.transpose(im*2-1, (2, 0, 1)).astype(np.float32)
im = np.transpose(im * 2 - 1, (2, 0, 1)).astype(np.float32)
return im


Expand Down Expand Up @@ -88,25 +87,25 @@ def fft2(K, size, dtype):

def laplacian_param(size, dtype):
w, h = size
K = np.zeros((2*w, 2*h)).astype(dtype)
K = np.zeros((2 * w, 2 * h)).astype(dtype)

laplacian_k = np.array([[0, -1, 0], [-1, 4, -1], [0, -1, 0]])
kw, kh = laplacian_k.shape
K[:kw, :kh] = laplacian_k
K = np.roll(K, -(kw//2), axis=0)
K = np.roll(K, -(kh//2), axis=1)

K = np.roll(K, -(kw // 2), axis=0)
K = np.roll(K, -(kh // 2), axis=1)

return fft2(K, size, dtype)


def gaussian_param(size, dtype, sigma):
w, h = size
K = np.zeros((2*w, 2*h)).astype(dtype)
K = np.zeros((2 * w, 2 * h)).astype(dtype)

K[1, 1] = 1
K[:3, :3] = gaussian(K[:3, :3], sigma)

K = np.roll(K, -1, axis=0)
K = np.roll(K, -1, axis=1)

Expand All @@ -126,41 +125,41 @@ def gaussian_poisson_editing(X, param_l, param_g, color_weight=1, eps=1e-12):
Fv = (X[:, :, :, 2] + np.roll(X[:, :, :, 4], -1, axis=0)) / 2
L = np.roll(Fh, 1, axis=1) + np.roll(Fv, 1, axis=0) - Fh - Fv

param = param_l + color_weight*param_g
param[(param >= 0) & (param < eps)] = eps
param[(param < 0) & (param > -eps)] = -eps
param = param_l + color_weight * param_g
param[(param >= 0) & (param < eps)] = eps
param[(param < 0) & (param > -eps)] = -eps

Y = np.zeros(X.shape[:3])
for i in range(3):
Xdct = dct2(X[:, :, i, 0])
Ydct = (dct2(L[:,:, i]) + color_weight*Xdct) / param
Y[:, :, i] = idct2(Ydct)
Xdct = dct2(X[:, :, i, 0])
Ydct = (dct2(L[:, :, i]) + color_weight * Xdct) / param
Y[:, :, i] = idct2(Ydct)
return Y


def run_gp_editing(src_im, dst_im, mask_im, gan_im, color_weight, sigma, gradient_kernel='normal'):
dst_feature = gradient_feature(dst_im, gan_im, gradient_kernel)
src_feature = gradient_feature(src_im, gan_im, gradient_kernel)
feature = dst_feature*(1-mask_im) + src_feature*mask_im
feature = dst_feature * (1 - mask_im) + src_feature * mask_im

size, dtype = feature.shape[:2], feature.dtype
param_l = laplacian_param(size, dtype)
param_g = gaussian_param(size, dtype, sigma)
gan_im = gaussian_poisson_editing(feature, param_l, param_g, color_weight=color_weight)
gan_im = np.clip(gan_im, 0, 1)
gan_im = gaussian_poisson_editing(feature, param_l, param_g, color_weight=color_weight)
gan_im = np.clip(gan_im, 0, 1)

return gan_im


def laplacian_pyramid(im, max_level, image_size, smooth_sigma):
im_pyramid = [im]
diff_pyramid = []
for i in range(max_level-1, -1, -1):
for i in range(max_level - 1, -1, -1):
smoothed = gaussian(im_pyramid[-1], smooth_sigma, multichannel=True)
diff_pyramid.append(im_pyramid[-1] - smoothed)
smoothed = ndarray_resize(smoothed, (image_size * 2**i, image_size * 2**i))
smoothed = ndarray_resize(smoothed, (image_size * 2 ** i, image_size * 2 ** i))
im_pyramid.append(smoothed)

im_pyramid.reverse()
diff_pyramid.reverse()

Expand All @@ -187,37 +186,43 @@ def laplacian_pyramid(im, max_level, image_size, smooth_sigma):
nz: noise vector lendth
n_iteration: # of iterations for optimization
"""
def gp_gan(obj, bg, mask, G, image_size, gpu, color_weight=1, sigma=0.5, gradient_kernel='normal', smooth_sigma=1, supervised=True, nz=100, n_iteration=1000):


def gp_gan(obj, bg, mask, G, image_size, gpu, color_weight=1, sigma=0.5, gradient_kernel='normal', smooth_sigma=1,
supervised=True, nz=100, n_iteration=1000):
w_orig, h_orig, _ = obj.shape
############################ Gaussian-Poisson GAN Image Editing ###########################
# pyramid
max_level = int(math.ceil(np.log2(max(w_orig, h_orig) / image_size)))
obj_im_pyramid, _ = laplacian_pyramid(obj, max_level, image_size, smooth_sigma)
bg_im_pyramid, _ = laplacian_pyramid(bg, max_level, image_size, smooth_sigma)
bg_im_pyramid, _ = laplacian_pyramid(bg, max_level, image_size, smooth_sigma)

# init GAN image
mask_init = ndarray_resize(mask, (image_size, image_size), order=0)[:, :, np.newaxis]
copy_paste_init = obj_im_pyramid[0]*mask_init + bg_im_pyramid[0]*(1-mask_init)
copy_paste_init = obj_im_pyramid[0] * mask_init + bg_im_pyramid[0] * (1 - mask_init)
copy_paste_init_var = Variable(chainer.dataset.concat_examples([preprocess(copy_paste_init)], gpu))

if supervised:
gan_im_var = G(copy_paste_init_var)
else:
z_init = np.random.normal(size=(nz, 1, 1))
res = minimize(z_generate, z_init, args=(G, copy_paste_init_var, nz, gpu), method='L-BFGS-B', jac=True, options={'maxiter': n_iteration, 'disp':False})
res = minimize(z_generate, z_init, args=(G, copy_paste_init_var, nz, gpu), method='L-BFGS-B', jac=True,
options={'maxiter': n_iteration, 'disp': False})
z = np.reshape(res.x, (nz, 1, 1)).astype(np.float32)
gan_im_var = G(Variable(chainer.dataset.concat_examples([z], gpu)))
gan_im = np.clip(np.transpose((np.squeeze(cuda.to_cpu(gan_im_var.data))+1)/2, (1, 2, 0)), 0, 1).astype(obj.dtype)
gan_im = np.clip(np.transpose((np.squeeze(cuda.to_cpu(gan_im_var.data)) + 1) / 2, (1, 2, 0)), 0, 1).astype(
obj.dtype)

# Start pyramid
for level in range(max_level+1):
for level in range(max_level + 1):
size = obj_im_pyramid[level].shape[:2]
mask_im = ndarray_resize(mask, size, order=0)[:, :, np.newaxis, np.newaxis]
if level != 0:
gan_im = ndarray_resize(gan_im, size)

gan_im = run_gp_editing(obj_im_pyramid[level], bg_im_pyramid[level], mask_im, gan_im, color_weight, sigma, gradient_kernel)

gan_im = np.clip(gan_im*255, 0, 255).astype(np.uint8)
gan_im = run_gp_editing(obj_im_pyramid[level], bg_im_pyramid[level], mask_im, gan_im, color_weight, sigma,
gradient_kernel)

gan_im = np.clip(gan_im * 255, 0, 255).astype(np.uint8)

return gan_im
19 changes: 9 additions & 10 deletions model.py
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import chainer
import chainer.links as L
import chainer.functions as F

import chainer.links as L
from chainer import cuda


Expand Down Expand Up @@ -41,7 +40,7 @@ def __call__(self, x):

class DCGAN_G(chainer.ChainList):
def __init__(self, isize, nc, ngf, conv_init=None, bn_init=None):
cngf, tisize = ngf//2, 4
cngf, tisize = ngf // 2, 4
while tisize != isize:
cngf = cngf * 2
tisize = tisize * 2
Expand All @@ -52,9 +51,9 @@ def __init__(self, isize, nc, ngf, conv_init=None, bn_init=None):
layers.append(L.BatchNormalization(cngf, initial_gamma=bn_init))
layers.append(ReLU())
csize, cndf = 4, cngf
while csize < isize//2:
layers.append(L.Deconvolution2D(None, cngf//2, ksize=4, stride=2, pad=1, initialW=conv_init, nobias=True))
layers.append(L.BatchNormalization(cngf//2, initial_gamma=bn_init))
while csize < isize // 2:
layers.append(L.Deconvolution2D(None, cngf // 2, ksize=4, stride=2, pad=1, initialW=conv_init, nobias=True))
layers.append(L.BatchNormalization(cngf // 2, initial_gamma=bn_init))
layers.append(ReLU())
cngf = cngf // 2
csize = csize * 2
Expand Down Expand Up @@ -97,17 +96,17 @@ def encode(self, x):
return x

def __call__(self, x):
x = self.encode(x)
x = self.encode(x)
x = F.sum(x, axis=0) / x.shape[0]
return F.squeeze(x)


class EncoderDecoder(chainer.Chain):
def __init__(self, nef, ngf, nc, nBottleneck, image_size=64, conv_init=None, bn_init=None):
super(EncoderDecoder, self).__init__(
encoder = DCGAN_D(image_size, nef, nBottleneck, conv_init, bn_init),
bn = L.BatchNormalization(nBottleneck, initial_gamma=bn_init),
decoder = DCGAN_G(image_size, nc, ngf, conv_init, bn_init)
encoder=DCGAN_D(image_size, nef, nBottleneck, conv_init, bn_init),
bn=L.BatchNormalization(nBottleneck, initial_gamma=bn_init),
decoder=DCGAN_G(image_size, nc, ngf, conv_init, bn_init)
)

def encode(self, x):
Expand Down
10 changes: 5 additions & 5 deletions requirements/test/environment.yml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name: gp-gan
dependencies:
- python=3.5
- pip=18.0
- chainer=6.3.0
- cupy=6.3.0
- scikit-image=0.15.0
- python=3.5
- pip=18.0
- chainer=6.3.0
- cupy=6.3.0
- scikit-image=0.15.0
Loading

0 comments on commit e166651

Please sign in to comment.