def add_y_to_ring_and_get_imshow(self, data):
ret = self.add_y_to_ring_and_get(data)
img = np.zeros(
(data.shape[1], self.image_height, self.image_width * 2),
dtype=dnn.dtype)
img[:, :, :self.image_width] = dnn.to_cpu(data[0])
img[:, :, self.image_width:] = dnn.to_cpu(ret[0])
cv2.imshow("y", dnn_to_bgr(img))
return ret
def run_test():
global xp
# 呼び出しを簡単にするため Node のメソッドとして組み込む
dnn.Node.cbr = cbr
dnn.Node.resblk = resblk
dnn.startup(0) # 0番のGPU使用
xp = dnn.xp
batch_size = 1
image_width = 640
image_height = 360
# モデル構築
m = CycleGan(640, 360, 3, 3, batch_size, 'train')
# 学習済みデータがあれば読み込む
print("Loading trained data...")
dnn.load_if_exists('cycle_gan', m)
print("Done.")
m = dnn.to_xp(m)
# データ用意
cap = cv2.VideoCapture("c:/work/ds3tamanegi.mp4")
cap.set(cv2.CAP_PROP_POS_FRAMES, 14000)
# バッチサイズ分のCPU上でのメモリ領域、一旦ここに展開してから to_xp すると無駄が少ない
xcpu = np.zeros((batch_size, 3, image_height, image_width),
dtype=dnn.dtype)
# 動画再生ループ
request_quit = False
while not request_quit:
success, image = cap.read()
if success:
image = cv2.resize(image, (image_width, image_height),
interpolation=cv2.INTER_AREA)
xcpu[0, :, :, :] = bgr_to_dnn(image)
y = m.test(dnn.to_xp(xcpu))
y = dnn.to_cpu(y.data[0])
image = dnn_to_bgr(y)
cv2.imshow("CycleGAN", image)
# OpenCVウィンドウアクティブにしてESCキーで中断できるようにしておく
k = cv2.waitKey(1)
if k == 27:
request_quit = True
break
def run_train():
global xp
# 呼び出しを簡単にするため Node のメソッドとして組み込む
dnn.Node.cbr = cbr
dnn.Node.resblk = resblk
dnn.startup(0) # 0番のGPU使用
xp = dnn.xp
batch_size = 3
image_width = 640
image_height = 360
image_size = (image_height, image_width)
# モデル構築
m = CycleGan(640, 360, 3, 3, batch_size, 'train')
# 学習済みデータがあれば読み込む
print("Loading trained data...")
# dnn.load_if_exists('cyclegan_data/cycle_gan', m)
# dnn.save('cyclegan_data2/cycle_gan', m)
# dnn.save_dict_to_hdf5('cycle_gan.h5', m.state_dict())
d = dnn.load_dict_from_hdf5('cycle_gan.h5')
m.load_state_dict(d)
# dnn.save('dict/cycle_gan', m)
print("Done.")
m = dnn.to_xp(m)
# バッチサイズ分のCPU上でのメモリ領域、一旦ここに展開してから to_xp すると無駄が少ない
xcpu = np.zeros((batch_size, 3, image_height, image_width),
dtype=dnn.dtype)
tcpu = np.zeros((batch_size, 3, image_height, image_width),
dtype=dnn.dtype)
# データ用意
x_caps = []
t_caps = []
for b in range(batch_size):
cap = cv2.VideoCapture("c:/work/Fortnite.mp4")
x_caps.append(cap)
cap = cv2.VideoCapture("c:/work/DarkSouls3Full.mp4")
t_caps.append(cap)
x_start_frame = 0
x_end_frame = int(x_caps[0].get(cv2.CAP_PROP_FRAME_COUNT))
t_start_frame = 0
t_end_frame = int(t_caps[0].get(cv2.CAP_PROP_FRAME_COUNT))
frame_step = 100
# 学習ループ
request_quit = False
iter_count = 0
while True:
if request_quit:
break
# バッチセットアップ
for i in range(batch_size):
x_cap = x_caps[i]
x_cap.set(cv2.CAP_PROP_POS_FRAMES,
random.randint(x_start_frame, x_end_frame))
success, image = x_cap.read()
if success:
image = cv2.resize(image, (image_width, image_height),
interpolation=cv2.INTER_AREA)
# image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
xcpu[i, :, :, :] = bgr_to_dnn(image)
t_cap = t_caps[i]
t_cap.set(cv2.CAP_PROP_POS_FRAMES,
random.randint(t_start_frame, t_end_frame))
success, image = t_cap.read()
if success:
image = cv2.resize(image, (image_width, image_height),
interpolation=cv2.INTER_AREA)
tcpu[i, :, :, :] = bgr_to_dnn(image)
# 学習実行
m.train(dnn.to_xp(xcpu), dnn.to_xp(tcpu), iter_count % 10 == 0,
iter_count % 10 == 0)
iter_count += 1
# OpenCVウィンドウアクティブにしてESCキーで中断できるようにしておく
k = cv2.waitKey(1)
if k == 27:
request_quit = True
break
def train(self, x, t, saveEval, showEval):
"""入力データと教師データを用いて学習を実行する.
# Args:
x: 入力データ. ※dnn.to_xp() で変換済みでなければならない
t: 教師データ. ※dnn.to_xp() で変換済みでなければならない
saveEval: 評価用画像を所定のディレクトリに保存するかどうか.
showEval: 評価用画像を表示するかどうか.
"""
self.iter += 1
x = x if isinstance(x, Variable) else Variable(x)
y = t if isinstance(t, Variable) else Variable(t)
x_y = self.gen_g(x)
x_y_copy = Variable(self.add_x_to_ring_and_get(x_y.data))
x_y_x = self.gen_f(x_y)
y_x = self.gen_f(y)
y_x_copy = Variable(self.add_y_to_ring_and_get(y_x.data))
y_x_y = self.gen_g(y_x)
if self.learning_rate_anneal > 0 and self.iter % self.learning_rate_anneal_interval == 0:
if self.gen_g.optimizer.alpha > self.learning_rate_anneal:
self.gen_g.optimizer.alpha -= self.learning_rate_anneal
if self.gen_f.optimizer.alpha > self.learning_rate_anneal:
self.gen_f.optimizer.alpha -= self.learning_rate_anneal
if self.dis_x.optimizer.alpha > self.learning_rate_anneal:
self.dis_x.optimizer.alpha -= self.learning_rate_anneal
if self.dis_y.optimizer.alpha > self.learning_rate_anneal:
self.dis_y.optimizer.alpha -= self.learning_rate_anneal
self.gen_g.zerograds()
self.gen_f.zerograds()
self.dis_x.zerograds()
self.dis_y.zerograds()
loss_dis_y_fake = loss_func_adv_dis_fake(self.dis_y(x_y_copy))
loss_dis_y_real = loss_func_adv_dis_real(self.dis_y(y))
loss_dis_y = loss_dis_y_fake + loss_dis_y_real
# chainer.report({'loss': loss_dis_y}, self.dis_y)
print('loss', loss_dis_y.data)
loss_dis_x_fake = loss_func_adv_dis_fake(self.dis_x(y_x_copy))
loss_dis_x_real = loss_func_adv_dis_real(self.dis_x(x))
loss_dis_x = loss_dis_x_fake + loss_dis_x_real
# chainer.report({'loss': loss_dis_x}, self.dis_x)
print('loss', loss_dis_x.data)
loss_dis_y.backward()
loss_dis_x.backward()
self.dis_y.optimizer.update()
self.dis_x.optimizer.update()
loss_gen_g_adv = loss_func_adv_gen(self.dis_y(x_y))
loss_gen_f_adv = loss_func_adv_gen(self.dis_x(y_x))
loss_cycle_x = self.lambda1 * loss_func_rec_l1(x_y_x, x)
loss_cycle_y = self.lambda1 * loss_func_rec_l1(y_x_y, y)
loss_gen = self.lambda2 * loss_gen_g_adv + self.lambda2 * loss_gen_f_adv + loss_cycle_x + loss_cycle_y
loss_gen.backward()
self.gen_f.optimizer.update()
self.gen_g.optimizer.update()
# chainer.report({'loss_rec': loss_cycle_y}, self.gen_g)
# chainer.report({'loss_rec': loss_cycle_x}, self.gen_f)
# chainer.report({'loss_adv': loss_gen_g_adv}, self.gen_g)
# chainer.report({'loss_adv': loss_gen_f_adv}, self.gen_f)
print('loss_rec', loss_cycle_y.data)
print('loss_rec', loss_cycle_x.data)
print('loss_adv', loss_gen_g_adv.data)
print('loss_adv', loss_gen_f_adv.data)
if saveEval or showEval:
w_in = self.image_width
h_in = self.image_height
img = np.zeros((x.data.shape[1], h_in * 2, w_in * 3),
dtype=dnn.dtype)
cells = [
dnn.to_cpu(x.data[0]),
dnn.to_cpu(x_y.data[0]),
dnn.to_cpu(x_y_x.data[0]),
dnn.to_cpu(y.data[0]),
dnn.to_cpu(y_x.data[0]),
dnn.to_cpu(y_x_y.data[0])
]
i = 0
for r in range(2):
rs = r * h_in
for c in range(3):
cs = c * w_in
img[:, rs:rs + h_in, cs:cs + w_in] = cells[i]
i += 1
img = dnn_to_bgr(img)
# if saveEval:
# f = os.path.normpath(os.path.join(dnn.GetEvalDataDir(), "eval.png"))
# cv2.imwrite(f, img)
# uploader.Upload(f)
if showEval:
cv2.imshow("CycleGAN", img)
return
def __init__(self,
image_width,
image_height,
in_ch,
out_ch,
batch_size,
mode,
learning_rate_g=0.0002,
learning_rate_d=0.0002,
lambda1=10.0,
lambda2=3.0,
learning_rate_anneal=0.0,
learning_rate_anneal_interval=1000):
"""モデル初期化.
Args:
image_width: 入力画像幅(px).
image_height: 入力画像高さ(px).
in_ch: 入力CH数.
out_ch: 出力CH数.
batch_size: バッチサイズ.
mode: モード、次のうちのどれか train / test.
args: コマンドライン引数の内モデル作成に必要な部分のみ抜き出したリスト.
"""
if mode == "train":
# モデル構築
self.gen_g = resblk_9(in_ch, out_ch)
self.gen_f = resblk_9(in_ch, out_ch)
self.dis_x = discriminator(in_ch)
self.dis_y = discriminator(in_ch)
# 推論の関数作成
self.build('cyclegan.prediction')
# ハイパーパラメータセットアップ
self.lambda1 = lambda1
self.lambda2 = lambda2
self.learning_rate_anneal = learning_rate_anneal
self.learning_rate_anneal_interval = learning_rate_anneal_interval
self.image_width = image_width
self.image_height = image_height
self.iter = 0
self.max_buffer_size = 50
# gen_g の出力を蓄えておくリングバッファ
self.ring_x = np.zeros((self.max_buffer_size, batch_size, in_ch,
self.image_height, self.image_width),
dtype=dnn.dtype)
self.ring_x_len = 0 # バッファ内有効要素数
self.ring_x_next = 0 # 次回バッファに書き込み時のインデックス
# gen_f の出力を蓄えておくリングバッファ
self.ring_y = np.zeros((self.max_buffer_size, batch_size, in_ch,
self.image_height, self.image_width),
dtype=dnn.dtype)
self.ring_y_len = 0 # バッファ内有効要素数
self.ring_y_next = 0 # 次回バッファに書き込み時のインデックス
# elif mode == "test":
# # テストモード用コマンドライン引数解析
# parser = argparse.ArgumentParser(description='Testing for CycleGAN.')
# parser.add_argument("-d", "--dataprefix", default="", help='Prefix for trained data file select.')
# parser.add_argument("-s", "--selectgen", default="gen_g", choices=['gen_g', 'gen_f'], help='Generator name "gen_g" or "gen_f".')
# parser.add_argument("-r", "--resblock", default="9", help='Number of Generator_ResBlock.', choices=['6', '9', '11', '12'])
# args = parser.parse_args(args)
# self.trainedDataPrefix = args.dataprefix
# # テスト用ニューラルネットワークセットアップ
# if args.selectgen == "gen_g":
# self.gen = generatorResBlock[args.resblock](in_ch, out_ch)
# self.chains["gen_g"] = self.gen
# elif args.selectgen == "gen_f":
# self.gen = generatorResBlock[args.resblock](in_ch, out_ch)
# self.chains["gen_f"] = self.gen
# else:
# print(args.selectgen, "このジェネレータは存在しません。")
# raise
else:
raise Exception("Not implemented mode. '{}'".format(mode))