def update(self):
com_optimizer = self.get_optimizer('com')
batch = self.get_iterator('main').next()
batchsize = len(batch)
mean_color = self.get_iterator('main').dataset.calc_mean_color()
cut_batch = GL_F.resize_and_cut(batch,
size_min=self.size_min,
size_max=self.size_max,
cut_size=self.cut_size)
masked_batch, mask_info = GL_F.make_mask(cut_batch,
size_min=self.mask_min,
size_max=self.mask_max,
paint=True,
color=mean_color,
stack=True)
# input_x:マスク前のカットされた画像 masked_x:マスクされ，マスク領域を示す4チャンネル目が追加された画像
# converterでCPU or GPU 用に変換
input_x = Variable(self.converter(cut_batch, self.device))
masked_x = Variable(self.converter(masked_batch, self.device))
# 補完器に入力
com = self.com
completed_x = com(masked_x)
# 出力の補完対象部分以外を入力で置き換える
replaced_x = GL_F.restore_input(input_x, completed_x, mask_info)
# update実行
com_optimizer.update(self.loss_com_MSE, com, replaced_x, input_x)
# イテレーション
self.iteration += 1
def update(self):
# optimizer
com_optimizer = self.get_optimizer('com')
dis_optimizer = self.get_optimizer('dis')
# ミニバッチをロード
batch = self.get_iterator('main').next()
batchsize = len(batch)
mean_color = self.get_iterator('main').dataset.calc_mean_color()
# バッチをリサイズ＆カットし，マスク付与
cut_batch = GL_F.resize_and_cut(batch,
size_min=self.size_min,
size_max=self.size_max,
cut_size=self.cut_size)
masked_batch, mask_info = GL_F.make_mask(cut_batch,
size_min=self.mask_min,
size_max=self.mask_max,
paint=True,
color=mean_color,
stack=True)
# input_x:マスク前のカットされた画像 masked_x:マスクされ，マスク領域を示す4チャンネル目が追加された画像
input_x = Variable(self.converter(cut_batch, self.device))
masked_x = Variable(self.converter(masked_batch, self.device))
# 補完・識別ネットワーク
com, dis = self.com, self.dis
# 補完ネットワークに入力
completed_x = com(masked_x)
# 補完器の出力の補完対象部分以外を入力で置き換える
replaced_x = GL_F.restore_input(input_x, completed_x, mask_info)
# Local Discriminator に入力する部分を切り出す
local_replaced_x = GL_F.cut_local(replaced_x,
cut_size=self.local_size,
mask_info=mask_info)
local_real_x = GL_F.cut_local(input_x, cut_size=self.local_size)
# Discriminatorに入力
completed_y = dis(local_replaced_x, replaced_x)
real_y = dis(local_real_x, input_x)
# update実行
com_optimizer.update(self.loss_com, com, replaced_x, input_x,
completed_y)
dis_optimizer.update(self.loss_dis, dis, completed_y, real_y)
# イテレーション
self.iteration += 1
def out_sample_image(dataset,
com,
rows,
cols,
seed,
dst,
name,
device,
out_input=True,
out_masked=True,
mean_color=None,
resize_min=256,
resize_max=None,
cut_size=256,
mask_min=96,
mask_max=128):
np.random.seed(seed)
n_images = rows * cols
xp = com.xp
device = device
dataset = dataset
converter = convert.concat_examples
i = np.random.randint(0, len(dataset) - n_images)
batch = dataset[i:i + n_images]
if mean_color == None:
mean_color = dataset.calc_mean_color()
start_1 = time.time()
cut_batch = GL_F.resize_and_cut(batch,
size_min=resize_min,
size_max=resize_max,
cut_size=cut_size)
masked_batch, mask_info = GL_F.make_mask(cut_batch,
size_min=mask_min,
size_max=mask_max,
paint=True,
color=mean_color,
stack=True)
# デバイス (CPU or GPU)にあわせてvariableを作る
input_x = Variable(converter(cut_batch, device))
masked_x = Variable(converter(masked_batch, device))
start_2 = time.time()
# 補完器に入力
# with chainer.using_config('train', False): # v2以降の仕様のようだ
# completed_x = com(masked_x)
completed_x = com(masked_x)
# 補完対象部分以外を元の画像で置き換える
replaced_x = GL_F.restore_input(input_x, completed_x, mask_info)
x = chainer.cuda.to_cpu(replaced_x.data)
z = masked_batch[:, 0:3]
org = cut_batch
np.random.seed()
x = np.asarray(np.clip(x * 255, 0.0, 255.0), dtype=np.uint8)
_, _, H, W = x.shape
x = x.reshape((rows, cols, 3, H, W))
x = x.transpose(0, 3, 1, 4, 2)
x = x.reshape((rows * H, cols * W, 3))
end = time.time()
print("--Calculation time--")
print("time1:{} [sec]".format(end - start_1))
print("time2:{} [sec]".format(end - start_2))
# 保存
if dst == '': dst = '.'
preview_dir = '{}'.format(dst)
x_path = preview_dir +\
'/completed_image_{}.png'.format(name)
if not os.path.exists(preview_dir):
os.makedirs(preview_dir)
Image.fromarray(x).save(x_path)
print('save completed image: {}'.format(x_path))
# 補完前の画像を出力
if out_input:
org = np.asarray(np.clip(org * 255, 0.0, 255.0), dtype=np.uint8)
org = org.reshape((rows, cols, 3, H, W))
org = org.transpose(0, 3, 1, 4, 2)
org = org.reshape((rows * H, cols * W, 3))
org_path = preview_dir +\
'/input_image_{}.png'.format(name)
Image.fromarray(org).save(org_path)
print('save original image: {}'.format(org_path))
# マスクを付与した画像を出力
if out_masked:
z = np.asarray(np.clip(z * 255, 0.0, 255.0), dtype=np.uint8)
z = z.reshape((rows, cols, 3, H, W))
z = z.transpose(0, 3, 1, 4, 2)
z = z.reshape((rows * H, cols * W, 3))
z_path = preview_dir +\
'/masked_image_{}.png'.format(name)
Image.fromarray(z).save(z_path)
print('save masked image: {}'.format(z_path))
def make_image(trainer):
np.random.seed(seed)
n_images = rows * cols
xp = com.xp
converter = trainer.updater.converter
device = trainer.updater.device
dataset = trainer.updater.get_iterator('main').dataset
i = np.random.randint(0, len(dataset) - n_images)
batch = dataset[i:i + n_images]
mean_color = dataset.calc_mean_color()
cut_batch = GL_F.resize_and_cut(batch,
size_min=resize_min,
size_max=resize_max,
cut_size=cut_size)
masked_batch, mask_info = GL_F.make_mask(cut_batch,
size_min=mask_min,
size_max=mask_max,
paint=True,
color=mean_color,
stack=True)
# デバイス (CPU or GPU)にあわせてvariableを作る
input_x = Variable(converter(cut_batch, device))
masked_x = Variable(converter(masked_batch, device))
# 補完器に入力
# with chainer.using_config('train', False): # v2以降の仕様のようだ
# completed_x = com(masked_x)
completed_x = com(masked_x)
# 補完対象部分以外を元の画像で置き換える
replaced_x = GL_F.restore_input(input_x, completed_x, mask_info)
x = chainer.cuda.to_cpu(replaced_x.data)
z = masked_batch[:, 0:3]
org = cut_batch
np.random.seed()
x = np.asarray(np.clip(x * 255, 0.0, 255.0), dtype=np.uint8)
z = np.asarray(np.clip(z * 255, 0.0, 255.0), dtype=np.uint8)
org = np.asarray(np.clip(org * 255, 0.0, 255.0), dtype=np.uint8)
_, _, H, W = x.shape
x = x.reshape((rows, cols, 3, H, W))
x = x.transpose(0, 3, 1, 4, 2)
x = x.reshape((rows * H, cols * W, 3))
z = z.reshape((rows, cols, 3, H, W))
z = z.transpose(0, 3, 1, 4, 2)
z = z.reshape((rows * H, cols * W, 3))
org = org.reshape((rows, cols, 3, H, W))
org = org.transpose(0, 3, 1, 4, 2)
org = org.reshape((rows * H, cols * W, 3))
preview_dir = '{}/preview'.format(dst)
x_path = preview_dir +\
'/completed_image{:0>8}.png'.format(trainer.updater.iteration)
z_path = preview_dir +\
'/masked_image_{:0>8}.png'.format(trainer.updater.iteration)
org_path = preview_dir +\
'/input_image_{:0>8}.png'.format(trainer.updater.iteration)
if not os.path.exists(preview_dir):
os.makedirs(preview_dir)
Image.fromarray(x).save(x_path)
Image.fromarray(z).save(z_path)
Image.fromarray(org).save(org_path)