def main():
place = paddle.set_device(FLAGS.device)
paddle.disable_static(place) if FLAGS.dynamic else None
im_shape = [-1, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
# Generators
g_AB = Generator()
g_BA = Generator()
g = paddle.Model(
GeneratorCombine(
g_AB, g_BA, is_train=False),
inputs=[input_A, input_B])
g.prepare()
g.load(FLAGS.init_model, skip_mismatch=True, reset_optimizer=True)
if not os.path.exists(FLAGS.output):
os.makedirs(FLAGS.output)
test_data_A = data.TestDataA()
test_data_B = data.TestDataB()
for i in range(len(test_data_A)):
data_A, A_name = test_data_A[i]
data_B, B_name = test_data_B[i]
data_A = np.array(data_A).astype("float32")
data_B = np.array(data_B).astype("float32")
fake_A, fake_B, cyc_A, cyc_B = g.test_batch([data_A, data_B])
datas = [fake_A, fake_B, cyc_A, cyc_B, data_A, data_B]
odatas = []
for o in datas:
d = np.squeeze(o[0]).transpose([1, 2, 0])
im = ((d + 1) * 127.5).astype(np.uint8)
odatas.append(im)
imsave(FLAGS.output + "/fakeA_" + B_name, odatas[0])
imsave(FLAGS.output + "/fakeB_" + A_name, odatas[1])
imsave(FLAGS.output + "/cycA_" + A_name, odatas[2])
imsave(FLAGS.output + "/cycB_" + B_name, odatas[3])
imsave(FLAGS.output + "/inputA_" + A_name, odatas[4])
imsave(FLAGS.output + "/inputB_" + B_name, odatas[5])
def main():
place = paddle.set_device(FLAGS.device)
paddle.disable_static(place) if FLAGS.dynamic else None
im_shape = [-1, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
# Generators
g_AB = Generator()
g_BA = Generator()
g = paddle.Model(GeneratorCombine(g_AB, g_BA, is_train=False),
inputs=[input_A, input_B])
g.prepare()
g.load(FLAGS.init_model, skip_mismatch=True, reset_optimizer=True)
out_path = FLAGS.output + "/single"
if not os.path.exists(out_path):
os.makedirs(out_path)
for f in glob.glob(FLAGS.input):
image_name = os.path.basename(f)
image = Image.open(f).convert('RGB')
image = image.resize((256, 256), Image.BICUBIC)
image = np.array(image) / 127.5 - 1
image = image[:, :, 0:3].astype("float32")
data = image.transpose([2, 0, 1])[np.newaxis, :]
if FLAGS.input_style == "A":
_, fake, _, _ = g.test_batch([data, data])
if FLAGS.input_style == "B":
fake, _, _, _ = g.test_batch([data, data])
fake = np.squeeze(fake[0]).transpose([1, 2, 0])
opath = "{}/fake{}{}".format(out_path, FLAGS.input_style, image_name)
imsave(opath, ((fake + 1) * 127.5).astype(np.uint8))
print("transfer {} to {}".format(f, opath))
def cyclegan(pretrained=None, device='cpu'):
""" # This docstring shows up in hub.help()
CycleGan model
pretrained (string): kwargs, load pretrained weights into the model (img2vangogh)
**kwargs
device (str): 'cuda' or 'cpu'
"""
net = Generator(3, 3)
if device == 'cuda':
net.cuda()
net.trained_models_list = [k for k in model_urls.keys()]
if pretrained is not None:
net.load_state_dict(
torch.hub.load_state_dict_from_url(
model_urls[pretrained],
progress=True,
map_location=torch.device(device)))
return net
def main():
place = set_device(FLAGS.device)
fluid.enable_dygraph(place) if FLAGS.dynamic else None
# Generators
g_AB = Generator()
g_BA = Generator()
# Discriminators
d_A = Discriminator()
d_B = Discriminator()
g = GeneratorCombine(g_AB, g_BA, d_A, d_B)
da_params = d_A.parameters()
db_params = d_B.parameters()
g_params = g_AB.parameters() + g_BA.parameters()
da_optimizer = opt(da_params)
db_optimizer = opt(db_params)
g_optimizer = opt(g_params)
im_shape = [None, 3, 256, 256]
input_A = Input(im_shape, 'float32', 'input_A')
input_B = Input(im_shape, 'float32', 'input_B')
fake_A = Input(im_shape, 'float32', 'fake_A')
fake_B = Input(im_shape, 'float32', 'fake_B')
g_AB.prepare(inputs=[input_A], device=FLAGS.device)
g_BA.prepare(inputs=[input_B], device=FLAGS.device)
g.prepare(g_optimizer,
GLoss(),
inputs=[input_A, input_B],
device=FLAGS.device)
d_A.prepare(da_optimizer,
DLoss(),
inputs=[input_B, fake_B],
device=FLAGS.device)
d_B.prepare(db_optimizer,
DLoss(),
inputs=[input_A, fake_A],
device=FLAGS.device)
if FLAGS.resume:
g.load(FLAGS.resume)
loader_A = paddle.io.DataLoader(data.DataA(),
places=place,
shuffle=True,
return_list=True,
batch_size=FLAGS.batch_size)
loader_B = paddle.io.DataLoader(data.DataB(),
places=place,
shuffle=True,
return_list=True,
batch_size=FLAGS.batch_size)
A_pool = data.ImagePool()
B_pool = data.ImagePool()
for epoch in range(FLAGS.epoch):
for i, (data_A, data_B) in enumerate(zip(loader_A, loader_B)):
data_A = data_A[0][0] if not FLAGS.dynamic else data_A[0]
data_B = data_B[0][0] if not FLAGS.dynamic else data_B[0]
start = time.time()
fake_B = g_AB.test_batch(data_A)[0]
fake_A = g_BA.test_batch(data_B)[0]
g_loss = g.train_batch([data_A, data_B])[0]
fake_pb = B_pool.get(fake_B)
da_loss = d_A.train_batch([data_B, fake_pb])[0]
fake_pa = A_pool.get(fake_A)
db_loss = d_B.train_batch([data_A, fake_pa])[0]
t = time.time() - start
if i % 20 == 0:
print("epoch: {} | step: {:3d} | g_loss: {:.4f} | " \
"da_loss: {:.4f} | db_loss: {:.4f} | s/step {:.4f}".
format(epoch, i, g_loss[0], da_loss[0], db_loss[0], t))
g.save('{}/{}'.format(FLAGS.checkpoint_path, epoch))
import torch
import onnx
# fix for python 3.6
try:
import cyclegan
except:
import sys
sys.path.insert(0, './')
from cyclegan import Generator
model_path = '/media/mint/Barracuda/Models/cyclegan/cezanne/downloaded/netG_B2A.pth'
netG_B2A = Generator(3, 3)
netG_B2A.load_state_dict(torch.load(model_path, map_location=(torch.device('cpu'))))
print('loaded...')
dummy_input = torch.randn(1, 3, 224, 224, device='cpu')
input_names = ["input"]
output_names = ["output"]
torch.onnx.export(netG_B2A,
dummy_input,
"static/cezanne.onnx",
verbose=True,
input_names=input_names,
output_names=output_names,
opset_version=10)
# Test
parser.add_argument('--output-folder',
type=str,
default='/media/mint/Barracuda/Models/cyclegan/vangog',
help='path to store the models')
opt = parser.parse_args()
print(opt)
if torch.cuda.is_available() and not opt.cuda:
print(
"WARNING: You have a CUDA device, so you should probably run with --cuda"
)
###### Definition of variables ######
# Networks
netG_A2B = Generator(opt.input_nc, opt.output_nc)
netG_B2A = Generator(opt.output_nc, opt.input_nc)
netD_A = Discriminator(opt.input_nc)
netD_B = Discriminator(opt.output_nc)
if opt.cuda:
netG_A2B.cuda()
netG_B2A.cuda()
netD_A.cuda()
netD_B.cuda()
netG_A2B.apply(weights_init_normal)
netG_B2A.apply(weights_init_normal)
netD_A.apply(weights_init_normal)
netD_B.apply(weights_init_normal)
from config import TOKEN
from cyclegan import Generator, device
from keyboards import menu, cancel_button, continue_button1, continue_button2
from nst import run_style_transfer, cnn, cnn_normalization_mean, cnn_normalization_std
from states import PhotoTransform
# Configure logging
logging.basicConfig(level=logging.INFO)
# Initialize bot and dispatcher
bot = Bot(TOKEN)
dp = Dispatcher(bot, storage=MemoryStorage())
# Импортируем модель класса Generator
gen = Generator().to(device)
# Выбор GPU или CPU
if torch.cuda.is_available():
map_location = lambda storage, loc: storage.cuda()
else:
map_location = 'cpu'
# Загрузка весов модели и устанавливаем инференс-режим
gen.load_state_dict(
torch.load('weights/netG_B2A.pth', map_location=map_location))
gen.eval()
""" БАЗОВЫЕ КОМАНДЫ /"""
@dp.message_handler(commands=['start'])