Skip to content

Commit

Permalink
Add files via upload
Browse files Browse the repository at this point in the history
  • Loading branch information
@HagopB
HagopB committed Nov 28, 2017
1 parent 5376373 commit e01fe1b
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 109 deletions.
31 changes: 18 additions & 13 deletions train.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,12 @@
parser = argparse.ArgumentParser()
parser.add_argument("--path_trainA", type=str, help="The path to the A style images")
parser.add_argument("--path_trainB", type=str, help="The path to the B style images")
parser.add_argument("--pic_dir", type=str, help="picture directory where to store all intermediate images")
parser.add_argument("--niter", type=int, help="Total number of iterations", default=100000)
parser.add_argument("--pic_dir", type=str, help="Picture directory where to store all intermediate images")
parser.add_argument("--lmbd", type=int, help="Lambada - weight of cycleloss", default=10)
parser.add_argument("--lmbd_feat", type=int, help="Lambda - weight of perception loss", default=0)
parser.add_argument("--niter", type=int, help="Total number of iterations", default=200)
parser.add_argument("--save_iter", type=int, help="Number of iterations before saving the model", default=250)
parser.add_argument("--cuda", type=str, help="cuda", default='2')
parser.add_argument("--cuda", type=str, help="cuda", default='1')
args = parser.parse_args()
print(args)

Expand All @@ -22,23 +24,26 @@
opt.batch_size = 1
opt.save_iter = args.save_iter
opt.niter = args.niter
opt.lmbd = 10
opt.lmbd = args.lmbd
opt.pic_dir = args.pic_dir
opt.idloss = 0.0
opt.lr = 0.0001
opt.lr = 0.0002
opt.d_iter = 1

if args.lmbd_feat != 0:
opt.perceptionloss = True
else:
opt.perceptionloss = False
opt.lmbd_feat = args.lmbd_feat

opt.__dict__.update(args.__dict__)
opt.summary()


cycleGAN = CycleGAN(opt)

IG_A = ImageGenerator(root=args.path_trainA,
resize=opt.resize,
crop=opt.crop)
IG_B = ImageGenerator(root=args.path_trainB,
resize=opt.resize,
crop=opt.crop)
IG = ImageGenerator(path_trainA=args.path_trainA,
path_trainB=args.path_trainB,
resize=opt.resize,
crop=opt.crop)

cycleGAN.fit(IG_A, IG_B)
cycleGAN.fit(IG)
199 changes: 103 additions & 96 deletions utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,25 +5,10 @@
import os
from keras.preprocessing.image import ImageDataGenerator
from keras.preprocessing import image

'''
*****************************************************************************
************************** BACKEND CHECK ********************************
*****************************************************************************
'''
def get_filter_dim():
'''
Theano uses `channels_first`: (batch, channels, height, width)
Tensorflow uses `channels_last`: (batch, height, width, channels)
In our case, tensorflow is used as backend
'''
data_format = K.image_data_format()
if data_format == 'channels_first':
return 1
elif data_format == 'channels_last':
return 3
else:
raise NotImplemented
from PIL import Image
import numpy as np
import glob
from random import randint, shuffle

'''
*****************************************************************************
Expand All @@ -46,6 +31,27 @@ def vis_grid(X, nh, nw, save_path=None):
imsave(save_path, img)
return img

def showG(A, B, path):
assert A.shape==B.shape
def G(fn_generate, X):
r = np.array([fn_generate([X[i:i+1]]) for i in range(X.shape[0])])
return r.swapaxes(0,1)[:,:,0]
rA = G(cycleA_generate, A)
rB = G(cycleB_generate, B)
arr = np.concatenate([A,B,rA[0],rB[0],rA[1],rB[1]])
saveX(arr, 3, path)

def saveX(X, path, rows=1):
imageSize=128
assert X.shape[0]%rows == 0
int_X = ( (X+1)/2*255).clip(0,255).astype('uint8')
int_X = int_X.reshape(-1,imageSize,imageSize, 3)
int_X = int_X.reshape(rows, -1, imageSize, imageSize,3).swapaxes(1,2).reshape(rows*imageSize,-1, 3)
img = Image.fromarray(int_X)
img.save(path)



'''
*****************************************************************************
************************** ImageGenerator ********************************
Expand All @@ -55,45 +61,59 @@ def vis_grid(X, nh, nw, save_path=None):
class ImageGenerator(object):
'''ImageGenerator'''
def __init__(self,
root,
path_trainA,
path_trainB,
resize=None,
crop=None,
flip=None):
self.img_list = os.listdir(root)
self.root = root
crop=None):
self.n_images_trainA_ = len(os.listdir(path_trainA))
self.n_images_trainB_ = len(os.listdir(path_trainB))
self.path_trainA = path_trainA
self.path_trainB = path_trainB
self.resize = resize
self.crop = crop
self.flip = flip

print('ImageGenerator from {} [{}]'.format(root, len(self.img_list)))

def __call__(self, bs):

def read_image(self, fn):
im = Image.open(fn).convert('RGB')
im = im.resize(self.resize, Image.BILINEAR )
arr = np.array(im)/255*2-1
w1, w2 = (self.resize[0] - self.crop[0])//2, (self.resize[0] + self.crop[0])//2
h1, h2 = w1,w2
img = arr[h1:h2, w1:w2, :]
if randint(0,1):
img=img[:,::-1]
return img

def minibatch(self, data, bs):
length = len(data)
epoch = i = 0
tmpsize = None
while True:
try:
imgs = []
for _ in range(bs):
img = imread(os.path.join(self.root, np.random.choice(self.img_list)))

if self.resize: img = imresize(img, self.resize)
if self.crop:
left = np.random.randint(0, img.shape[0]-self.crop[0])
top = np.random.randint(0, img.shape[1]-self.crop[1])
img = img[left:left+self.crop[0], top:top+self.crop[1]]
if self.flip:
if np.random.random() > 0.5:
img = img[:, ::-1, :]

imgs.append(img)

imgs = np.array(imgs)
if get_filter_dim() == 1:
imgs = imgs.transpose(0, 3, 1, 2)
size = tmpsize if tmpsize else bs
if i + size > length:
shuffle(data)
i = 0
epoch+=1
rtn = [self.read_image(data[j]) for j in range(i, i + size)]
i += size
tmpsize = yield epoch, np.float32(rtn)

def minibatchAB(self, dataA, dataB, bs):
batchA = self.minibatch(dataA, bs)
batchB = self.minibatch(dataB, bs)
tmpsize = None
while True:
ep1, A = batchA.send(tmpsize)
ep2, B = batchB.send(tmpsize)
tmpsize = yield max(ep1, ep2), A, B

def __call__(self, bs):
trainA = glob.glob('{}/*'.format(self.path_trainA))
trainB = glob.glob('{}/*'.format(self.path_trainB))

imgs = imgs/127.5-1
print('N images train A {} -- N images train B {}'.format(len(trainA), len(trainB)))

return imgs
except:
pass
return self.minibatchAB(trainA, trainB, bs)


'''
*****************************************************************************
Expand All @@ -107,62 +127,49 @@ class Option(object):
def __init__(self,
# from CycleGAN/options.lua
# data
DATA_ROOT = '', ## path to images (should have subfolders 'train', 'val', etc)
shapeA = (256,256,3),
shapeB = (256,256,3),
resize = (286,286),
crop = (256,256),
DATA_ROOT = '', # path to images (should have subfolders 'train', 'val', etc)
shapeA = (128,128,3), #(256,256,3),
shapeB = (128,128,3), #(256,256,3),
resize = (143,143), #(286,286),
crop = (128,128), #(256,256),

# net definition
which_model_netD = 'basic', ## selects model to use for netD
which_model_netG = 'resnet_6blocks', ## selects model to use for netG
use_lsgan = 1, ## if 1, use least square GAN, if 0, use vanilla GAN
ngf = 64, ## # of gen filters in first conv layer
ndf = 64, ## # of discrim filters in first conv layer
which_model_netD = 'basic', # selects model to use for netD
which_model_netG = 'unet_128', # selects model to use for netG
use_lsgan = 1, # if 1, use least square GAN, if 0, use vanilla GAN
perceptionloss = False, # wether to use CycleGan with perception loss
ngf = 64, # # of gen filters in first conv layer
ndf = 64, # of discrim filters in first conv layer
lmbd = 10.0,
idloss=1.0,

lmbd_feat = 1.0,
# optimizers
lr = 0.0002, ## initial learning rate for adam
beta1 = 0.5, ## momentum term of adam
lr = 0.0002, # initial learning rate for adam
beta1 = 0.5, # momentum term of adam

# training parameters
batch_size = 1, ## # images in batch
niter = 100, ## # of iter at starting learning rate
pool_size = 50, ## the size of image buffer that stores previously generated images
batch_size = 1, # images in batch
niter = 200, # of iter at starting learning rate
pool_size = 50, # the size of image buffer that stores previously generated images
save_iter = 50,
d_iter = 10,

# dirs
pic_dir = 'quickshots',

niter_decay = 100, ## # of iter to linearly decay learning rate to zero
ntrain = np.inf, ## # of examples per epoch. math.huge for full dataset
flip = 1, ## if flip the images for data argumentation
display_id = 10, ## display window id.
display_winsize = 256, ## display window size
display_freq = 25, ## display the current results every display_freq iterations
gpu = 1, ## gpu = 0 is CPU mode. gpu=X is GPU mode on GPU X
name = '', ## name of the experiment, should generally be passed on the command line
which_direction = 'AtoB', ## AtoB or BtoA
phase = 'train', ## train, val, test, etc
nThreads = 2, ## # threads for loading data
save_epoch_freq = 1, ## save a model every save_epoch_freq epochs (does not overwrite previously saved models)
save_latest_freq = 5000, ## save the latest model every latest_freq sgd iterations (overwrites the previous latest model)
print_freq = 50, ## print the debug information every print_freq iterations
save_display_freq = 2500, ## save the current display of results every save_display_freq_iterations
continue_train = 0, ## if continue training, load the latest model: 1: true, 0: false
serial_batches = 0, ## if 1, takes images in order to make batches, otherwise takes them randomly
checkpoints_dir = './checkpoints', ## models are saved here
cudnn = 1, ## set to 0 to not use cudnn
norm = 'instance', ## batch or instance normalization
n_layers_D = 3, ## only used if which_model_netD=='n_layers'
content_loss = 'pixel', ## content loss type: pixel, vgg
layer_name = 'pixel', ## layer used in content loss (e.g. relu4_2)
model = 'cycle_gan', ## which mode to run. 'cycle_gan', 'pix2pix', 'bigan', 'content_gan'
align_data = 0, ## if > 0, use the dataloader for where the images are aligned
resize_or_crop = 'resize_and_crop', ## resizing/cropping strategy
identity = 0, ## use identity mapping. Setting opt.identity other than 1 has an effect of scaling the weight of the identity mapping loss. For example, if the weight of the identity loss should be 10 times smaller than the weight of the reconstruction loss, please set opt.identity = 0.1
niter_decay = 100, # of iter to linearly decay learning rate to zero
ntrain = np.inf, # of examples per epoch. math.huge for full dataset
flip = 1, # if flip the images for data argumentation
display_id = 10, # display window id.
display_winsize = 128, # 256 if images are of shape (256, 256, 3)
display_freq = 25, # display the current results every display_freq iterations
gpu = 1, # gpu = 0 is CPU mode. gpu=X is GPU mode on GPU X
cuda=1,
name = '', # name of the experiment, should generally be passed on the command line
save_epoch_freq = 1, # save a model every save_epoch_freq epochs (does not overwrite previously saved models)
save_latest_freq = 5000, # save the latest model every latest_freq sgd iterations (overwrites the previous latest model)
print_freq = 50, # print the debug information every print_freq iterations
save_display_freq = 2500, # save the current display of results every save_display_freq_iterations
):
#assert shapeA[0:1] == crop
self.__dict__.update(locals())
Expand Down

0 comments on commit e01fe1b

Please sign in to comment.