def stylize(args):
device = torch.device("cuda" if args.cuda else "cpu")
style_model = TransformerNet()
state_dict = torch.load(args.model)
style_model.load_state_dict(state_dict)
style_model.to(device)
img_list = os.listdir(args.content_dir)
img_list.sort()
for img in tqdm(img_list):
img_path = args.content_dir + img
content_org = utils.load_image(img_path, scale=args.content_scale)
content_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))])
content_image = content_transform(content_org)
content_image = content_image.unsqueeze(0).to(device)
with torch.no_grad():
output = style_model(content_image).cpu()
output = output[0]
output = output.clone().clamp(0, 255).numpy()
output = output.transpose(1, 2, 0).astype("uint8")
output = Image.fromarray(output)
if args.keep_colors:
output = utils.original_colors(content_org, output)
output.save(args.output_dir + img)
def test():
image_path = './image_folder'
ckpt = './path_to_ckpt'
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = test_transform()
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(ckpt))
transformer.eval()
image = torch.Tensor(transform(Image.open(image_path))).to(device)
image = image.unsqueeze(0)
with torch.no_grad():
style_trans_img = denormalize(transformer(image)).cpu()
fn = image_path.split("/")[-1]
save_image(style_trans_img, f"images/outputs/stylized-{fn}")
def stylize():
net = TransformerNet()
net.load_state_dict(torch.load(STYLE_TRANSFORM_PATH))
net = net.to(device)
with torch.no_grad():
while (1):
torch.cuda.empty_cache()
print("Stylize Image~ Press Ctrl+C and Enter to close the program")
content_image_path = input("Enter the image path: ")
content_image = cv2.imread(content_image_path)
content_tensor = itot(content_image)
generated_tensor = net(content_tensor)
generated_image = ttoi(generated_tensor)
generated_image = cv2.cvtColor(generated_image, cv2.COLOR_BGR2RGB)
plt.imshow(generated_image)
plt.show()
def style_image(image_path, model_path):
image = Image.open(image_path)
width, height = image.size
alpha = image.convert('RGBA').split()[-1]
# @TODO - import the mean color...
mean_color = Image.new("RGB", image.size, (124, 116, 103))
rgb_image = image.convert('RGB')
rgb_image.paste(mean_color, mask=invert(alpha))
cuda_available = torch.cuda and torch.cuda.is_available()
device = torch.device("cuda" if cuda_available else "cpu")
model = torch.load(model_path, map_location=device)
image_filename = os.path.basename(image_path)
model_filename = os.path.basename(model_path)
model_name = os.path.splitext(model_filename)[0]
os.makedirs(f"images/outputs/{model_name}", exist_ok=True)
transform = style_transform()
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(model)
transformer.eval()
# Prepare input
image_tensor = Variable(transform(rgb_image)).to(device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
output_tensor = depro(transformer(image_tensor))
stylized_image = F.to_pil_image(output_tensor) \
.convert('RGBA') \
.crop((0, 0, width, height))
stylized_image.putalpha(alpha)
stylized_image.save(f"images/outputs/{model_name}/{image_filename}", 'PNG')
def transfer_img(usr_img_path, style_model_path, new_img_path):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(style_model_path))
transformer.eval()
# Prepare input
transform = style_transform()
image_tensor = transform(Image.open(usr_img_path)).to(device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
# Save image
save_image(stylized_image, new_img_path)
def exportStyleTransfer(image_path, style):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
styles = ['9', '32', '51', '55', '56', '58', '108', '140', '150', '153', '154', '155', '156']
checkpoint_model = '/home/KLTN_TheFaceOfArtFaceParsing/Updates/StyleTransfer/models/' + styles[style-1] + '_4000.pth'
transform = style_transform()
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(checkpoint_model))
transformer.eval()
# Prepare input
image_tensor = Variable(transform(Image.open(image_path))).to(device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
# Save image
save_image(stylized_image,'/home/KLTN_TheFaceOfArtFaceParsing/result.jpg')
def stylize(**kwargs):
opt.parse(**kwargs)
device = t.device('cuda') if opt.use_gpu else t.device('cpu')
# data
# 对单张图片进行加载验证
content_image = tv.datasets.folder.default_loader(opt.content_path)
# 这个transform更看不懂了，竟然不需要resize和centrocrop，这图片大小不一致，能送进去么？？？？
content_transform = tv.transforms.Compose([
tv.transforms.ToTensor(),
tv.transforms.Lambda(lambda x: x.mul(255))
])
content_image = content_transform(content_image)
content_image = content_image.unsqueeze(0).to(device).detach()
# model
style_model = TransformerNet().eval()
style_model.load_state_dict(t.load(opt.model_path, map_location=lambda _s, _: _s))
style.to(device)
# stylize
output = style_model(content_image) # 0-255~
output_data = (output.cpu().data[0]/255).clamp(min=0, max=1)
tv.utils.save_image(output_data,opt.result_path)
class Inference():
"""
Inference class that wrapped styleTransfer classes.
"""
def __init__(self):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.mosaic_WeightsPath = "Fast-Neural-Style-Transfer/weights/mosaic_10000.pth"
self.cuphead_WeightsPath = mospath + "Fast-Neural-Style-Transfer/weights/cuphead_10000.pth"
self.starry_night_WeightsPath = mospath + "Fast-Neural-Style-Transfer/weights/starry_night_10000.pth"
# instance of style Transfer model.
self.transformer = TransformerNet().to(self.device)
self.transform = style_transform()
def load_model(self, modelType):
"""
Args:
modelType: string that contains name of the model to switch to it.
"""
if modelType.lower() == "mosaic":
self.transformer.load_state_dict(torch.load(
self.mosaic_WeightsPath, map_location=torch.device('cpu')))
elif modelType.lower() == "cuphead":
self.transformer.load_state_dict(torch.load(
self.cuphead_WeightsPath, map_location=torch.device('cpu')))
else:
self.transformer.load_state_dict(torch.load(
self.starry_night_WeightsPath, map_location=torch.device('cpu')))
self.transformer.eval()
def predict(self, modelType, Base64Img):
img = base64.b64decode(Base64Img)
image = Image.open(BytesIO(img)).convert('RGB')
self.load_model(modelType)
image_tensor = Variable(self.transform(image)).to(self.device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = denormalize(self.transformer(image_tensor)).cpu()
imageBytes = save_image(
stylized_image, "/home/mostafax/Desktop/Style-Transfer-App/Fast-Neural-Style-Transfer/images/result.jpeg")
my_string = base64.b64encode(imageBytes)
return my_string
def main():
uploaded_file = st.file_uploader(
"Choose a picc", type=['jpg', 'png', 'webm', 'mp4', 'gif', 'jpeg'])
if uploaded_file is not None:
st.image(uploaded_file, width=200)
folder = os.path.abspath(os.getcwd())
folder = folder + '/models'
fnames = []
for basename in os.listdir(folder):
print(basename)
fname = os.path.join(folder, basename)
if fname.endswith('.pth'):
fnames.append(fname)
checkpoint = st.selectbox('Select a pretrained model', fnames)
# parser = argparse.ArgumentParser()
# parser.add_argument("--image_path", type=str, required=True, help="Path to image")
# parser.add_argument("--checkpoint_model", type=str, required=True, help="Path to checkpoint model")
# args = parser.parse_args()
# print(args)
os.makedirs("images/outputs", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = torch.device("cpu")
transform = style_transform()
try:
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(checkpoint))
transformer.eval()
# Prepare input
image_tensor = Variable(
transform(Image.open(uploaded_file).convert('RGB'))).to(device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
# colormaps = ['Accent', 'Accent_r', 'Blues', 'Blues_r', 'BrBG', 'BrBG_r', 'BuGn', 'BuGn_r', 'BuPu', 'BuPu_r', 'CMRmap', 'CMRmap_r', 'Dark2', 'Dark2_r', 'GnBu', 'GnBu_r', 'Greens', 'Greens_r', 'Greys', 'Greys_r', 'OrRd', 'OrRd_r', 'Oranges', 'Oranges_r', 'PRGn', 'PRGn_r', 'Paired', 'Paired_r', 'Pastel1', 'Pastel1_r', 'Pastel2', 'Pastel2_r', 'PiYG', 'PiYG_r', 'PuBu', 'PuBuGn', 'PuBuGn_r', 'PuBu_r', 'PuOr', 'PuOr_r', 'PuRd', 'PuRd_r', 'Purples', 'Purples_r', 'RdBu', 'RdBu_r', 'RdGy', 'RdGy_r', 'RdPu', 'RdPu_r', 'RdYlBu', 'RdYlBu_r', 'RdYlGn', 'RdYlGn_r', 'Reds', 'Reds_r', 'Set1', 'Set1_r', 'Set2', 'Set2_r', 'Set3', 'Set3_r', 'Spectral', 'Spectral_r', 'Wistia', 'Wistia_r', 'YlGn', 'YlGnBu', 'YlGnBu_r', 'YlGn_r', 'YlOrBr', 'YlOrBr_r', 'YlOrRd', 'YlOrRd_r', 'afmhot', 'afmhot_r', 'autumn', 'autumn_r', 'binary', 'binary_r', 'bone', 'bone_r', 'brg', 'brg_r', 'bwr', 'bwr_r', 'cividis', 'cividis_r', 'cool', 'cool_r', 'coolwarm', 'coolwarm_r', 'copper', 'copper_r', 'cubehelix', 'cubehelix_r', 'flag', 'flag_r', 'gist_earth', 'gist_earth_r', 'gist_gray', 'gist_gray_r', 'gist_heat', 'gist_heat_r', 'gist_ncar', 'gist_ncar_r', 'gist_rainbow', 'gist_rainbow_r', 'gist_stern', 'gist_stern_r', 'gist_yarg', 'gist_yarg_r', 'gnuplot', 'gnuplot2', 'gnuplot2_r', 'gnuplot_r', 'gray', 'gray_r', 'hot', 'hot_r', 'hsv', 'hsv_r', 'inferno', 'inferno_r', 'jet', 'jet_r', 'magma', 'magma_r', 'nipy_spectral', 'nipy_spectral_r', 'ocean', 'ocean_r', 'pink', 'pink_r', 'plasma', 'plasma_r', 'prism', 'prism_r', 'rainbow', 'rainbow_r', 'seismic', 'seismic_r', 'spring', 'spring_r', 'summer', 'summer_r', 'tab10', 'tab10_r', 'tab20', 'tab20_r', 'tab20b', 'tab20b_r', 'tab20c', 'tab20c_r', 'terrain', 'terrain_r', 'turbo', 'turbo_r', 'twilight', 'twilight_r', 'twilight_shifted', 'twilight_shifted_r', 'viridis', 'viridis_r', 'winter', 'winter_r']
# colormap = st.selectbox('Select a colormap', colormaps)
# plt.imshow(stylized_image.numpy()[0][0], cmap=colormap)
# plt.imshow(stylized_image.numpy()[0][0], cmap='gist_rainbow')
# img = np.squeeze(stylized_image)
# plt.imshow(img[0])
# plt.show()
# st.image(img)
# # Save image
fn = str(np.random.randint(0, 100)) + 'image.jpg'
save_image(stylized_image, f"images/outputs/stylized-{fn}")
st.image(f"images/outputs/stylized-{fn}")
except:
st.write('Choose an image')
device = get_device()
content_image = load_image(test_image)
content_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
content_image = content_transform(content_image)
content_image = content_image.unsqueeze(0).to(device)
with torch.no_grad():
style_model = TransformerNet()
ckpt_model_path = os.path.join(checkpoint_dir, checkpoint_file)
checkpoint = torch.load(ckpt_model_path, map_location=device)
# remove saved deprecated running_* keys in InstanceNorm from the checkpoint
for k in list(checkpoint.keys()):
if re.search(r'in\d+\.running_(mean|var)$', k): # in200.running_var or in200.running_mean
del checkpoint[k]
style_model.load_state_dict(checkpoint['model_state_dict'])
style_model.to(device)
output = style_model(content_image).cpu()
save_image(output_image, output[0])
print('Save image !!')
def get_transformer(checkpoint_path: str) -> torch.nn.Module:
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(checkpoint_path, map_location=device))
transformer.eval()
return transformer
def train(**kwargs):
# step1:config
opt.parse(**kwargs)
vis = Visualizer(opt.env)
device = t.device('cuda') if opt.use_gpu else t.device('cpu')
# step2:data
# dataloader, style_img
# 这次图片的处理和之前不一样，之前都是normalize,这次改成了lambda表达式乘以255，这种转化之后要给出一个合理的解释
# 图片共分为两种，一种是原图，一种是风格图片，在作者的代码里，原图用于训练，需要很多，风格图片需要一张，用于损失函数
transforms = T.Compose([
T.Resize(opt.image_size),
T.CenterCrop(opt.image_size),
T.ToTensor(),
T.Lambda(lambda x: x*255)
])
# 这次获取图片的方式和第七章一样，仍然是ImageFolder的方式，而不是dataset的方式
dataset = tv.datasets.ImageFolder(opt.data_root,transform=transforms)
dataloader = DataLoader(dataset,batch_size=opt.batch_size,shuffle=True,num_workers=opt.num_workers,drop_last=True)
style_img = get_style_data(opt.style_path) # 1*c*H*W
style_img = style_img.to(device)
vis.img('style_image',(style_img.data[0]*0.225+0.45).clamp(min=0,max=1)) # 个人觉得这个没必要，下次可以实验一下
# step3: model：Transformer_net 和 损失网络vgg16
# 整个模型分为两部分，一部分是转化模型TransformerNet，用于转化原始图片，一部分是损失模型Vgg16，用于评价损失函数，
# 在这里需要注意一下，Vgg16只是用于评价损失函数的，所以它的参数不参与反向传播，只有Transformer的参数参与反向传播，
# 也就意味着，我们只训练TransformerNet，只保存TransformerNet的参数，Vgg16的参数是在网络设计时就已经加载进去的。
# Vgg16是以验证model.eval()的方式在运行，表示其中涉及到pooling等层会发生改变
# 那模型什么时候开始model.eval()呢，之前是是val和test中就会这样设置，那么Vgg16的设置理由是什么？
# 这里加载模型的时候，作者使用了简单的map_location的记录方法，更轻巧一些
# 发现作者在写这些的时候越来越趋向方便的方式
# 在cuda的使用上，模型的cuda是直接使用的，而数据的cuda是在正式训练的时候才使用的，注意一下两者的区别
# 在第七章作者是通过两种方式实现网络分离的，一种是对于前面网络netg,进行 fake_img = netg(noises).detach(),使得非叶子节点变成一个类似不需要邱求导的叶子节点
# 第四章还需要重新看，
transformer_net = TransformerNet()
if opt.model_path:
transformer_net.load_state_dict(t.load(opt.model_path,map_location= lambda _s, _: _s))
transformer_net.to(device)
# step3： criterion and optimizer
optimizer = t.optim.Adam(transformer_net.parameters(),opt.lr)
# 此通过vgg16实现的，损失函数包含两个Gram矩阵和均方误差，所以，此外，我们还需要求Gram矩阵和均方误差
vgg16 = Vgg16().eval() # 待验证
vgg16.to(device)
# vgg的参数不需要倒数，但仍然需要反向传播
# 回头重新考虑一下detach和requires_grad的区别
for param in vgg16.parameters():
param.requires_grad = False
criterion = t.nn.MSELoss(reduce=True, size_average=True)
# step4: meter 损失统计
style_meter = meter.AverageValueMeter()
content_meter = meter.AverageValueMeter()
total_meter = meter.AverageValueMeter()
# step5.2：loss 补充
# 求style_image的gram矩阵
# gram_style:list [relu1_2,relu2_2,relu3_3,relu4_3] 每一个是b*c*c大小的tensor
with t.no_grad():
features = vgg16(style_img)
gram_style = [gram_matrix(feature) for feature in features]
# 损失网络 Vgg16
# step5： train
for epoch in range(opt.epoches):
style_meter.reset()
content_meter.reset()
# step5.1: train
for ii,(data,_) in tqdm(enumerate(dataloader)):
optimizer.zero_grad()
# 这里作者没有进行 Variable(),与之前不同
# pytorch 0.4.之后tensor和Variable不再严格区分，创建的tensor就是variable
# https://mp.weixin.qq.com/s?__biz=MzI0ODcxODk5OA==&mid=2247494701&idx=2&sn=ea8411d66038f172a2f553770adccbec&chksm=e99edfd4dee956c23c47c7bb97a31ee816eb3a0404466c1a57c12948d807c975053e38b18097&scene=21#wechat_redirect
data = data.to(device)
y = transformer_net(data)
# vgg对输入的图片需要进行归一化
data = normalize_batch(data)
y = normalize_batch(y)
feature_data = vgg16(data)
feature_y = vgg16(y)
# 疑问？？现在的feature是一个什么样子的向量？
# step5.2: loss:content loss and style loss
# content_loss
# 在这里和书上的讲的不一样，书上是relu3_3,代码用的是relu2_2
# https://blog.csdn.net/zhangxb35/article/details/72464152?utm_source=itdadao&utm_medium=referral
# 均方误差指的是一个像素点的损失，可以理解N*b*h*w个元素加起来，然后除以N*b*h*w
# 随机梯度下降法本身就是对batch内loss求平均后反向传播
content_loss = opt.content_weight*criterion(feature_y.relu2_2,feature_data.relu2_2)
# style loss
# style loss:relu1_2,relu2_2,relu3_3,relu3_4
# 此时需要求每一张图片的gram矩阵
style_loss = 0
# tensor也可以 for i in tensor:,此时只拆解外面一层的tensor
# ft_y:b*c*h*w, gm_s:1*c*h*w
for ft_y, gm_s in zip(feature_y, gram_style):
gram_y = gram_matrix(ft_y)
style_loss += criterion(gram_y, gm_s.expand_as(gram_y))
style_loss *= opt.style_weight
total_loss = content_loss + style_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
#import ipdb
#ipdb.set_trace()
# 获取tensor的值 tensor.item() tensor.tolist()
content_meter.add(content_loss.item())
style_meter.add(style_loss.item())
total_meter.add(total_loss.item())
# step5.3: visualize
if (ii+1)%opt.print_freq == 0 and opt.vis:
# 为什么总是以这种形式进行debug
if os.path.exists(opt.debug_file):
import ipdb
ipdb.set_trace()
vis.plot('content_loss',content_meter.value()[0])
vis.plot('style_loss',style_meter.value()[0])
vis.plot('total_loss',total_meter.value()[0])
# 因为现在data和y都已经经过了normalize，变成了-2~2，所以需要把它变回去0-1
vis.img('input',(data.data*0.225+0.45)[0].clamp(min=0,max=1))
vis.img('output',(y.data*0.225+0.45)[0].clamp(min=0,max=1))
# step 5.4 save and validate and visualize
if (epoch+1) % opt.save_every == 0:
t.save(transformer_net.state_dict(), 'checkpoints/%s_style.pth' % epoch)
# 保存图片的几种方法，第七章的是
# tv.utils.save_image(fix_fake_imgs,'%s/%s.png' % (opt.img_save_path, epoch),normalize=True, range=(-1,1))
# vis.save竟然没找到 我的神
vis.save([opt.env])
def index():
if request.method == 'POST':
f = request.files['image']
filename = f.filename
path = os.path.join(UPLOAD_FOLDER, filename)
f.save(path)
style_img = path
a = int(Image.open(style_img).size[0])
b = int(Image.open(style_img).size[1])
device = torch.device("cpu")
transform = style_transform()
transformer = TransformerNet().to(device)
model = request.form['style']
if model == 'mosaic':
transformer.load_state_dict(
torch.load("static/model/mosaic_10000.pth",
map_location=torch.device('cpu')))
transformer.eval()
filename = "mosaic{}".format(filename)
elif model == 'mona':
transformer.load_state_dict(
torch.load("static/model/mona_24000.pth",
map_location=torch.device('cpu')))
transformer.eval()
filename = "mona{}".format(filename)
elif model == 'starry':
transformer.load_state_dict(
torch.load("static/model/starry_night_10000.pth",
map_location=torch.device('cpu')))
transformer.eval()
filename = "starry{}".format(filename)
# Prepare input
if a * b < 800000:
image_tensor = Variable(
transform(Image.open(style_img).convert("RGB"))).to(device)
image_tensor = image_tensor.unsqueeze(0)
elif 800000 < a * b < 1960000:
image_tensor = Variable(
transform(
Image.open(style_img).convert("RGB").resize(
(int(a * 2 / 3), int(b * 2 / 3))))).to(device)
image_tensor = image_tensor.unsqueeze(0)
else:
image_tensor = Variable(
transform(
Image.open(style_img).convert("RGB").resize(
(int(a / 2), int(b / 2))))).to(device)
image_tensor = image_tensor.unsqueeze(0)
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
save_image(stylized_image,
"./static/predict/result_{}".format(filename))
# stylized_image.save("./static/predict/{}".format(filename))
#prediction pass to pipeline model
return render_template("index.html",
fileupload=True,
img_name="result_" + filename)
return render_template("index.html", fileupload=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mean, std = calculate_mean_and_std(args.dataset_path, args.image_size)
print(
"Please save the mean and standard deviation, this is needed when running inference."
)
print("Mean: ", mean, " STD: ", std)
# Create data loader for the training data
train_dataset = datasets.ImageFolder(
args.dataset_path, train_transform(args.image_size, mean, std))
dataloader = DataLoader(train_dataset, batch_size=args.batch_size)
# Defines networks
transformer = TransformerNet().to(device)
vgg = VGG16(requires_grad=False).to(device)
if args.checkpoint_model:
transformer.load_state_dict(
torch.load(args.checkpoint_model, device=device))
optimizer = Adam(transformer.parameters(), args.lr)
l2_loss = torch.nn.MSELoss().to(device)
# Load style image
style = get_preprocessor(args.style_size)(args.style_image).squeeze(0)
style = style.repeat(args.batch_size, 1, 1, 1).to(device)
# Extract style features
features_style = vgg(style)
gram_style = [gram_matrix(y) for y in features_style]
# Sample 8 images for visual evaluation of the model
image_samples = []
return db_op_wrapper
####################################################################################
# LOAD AND COMPILE MODEL
####################################################################################
size = 640, 640
transform = style_transform()
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# ohne gpu ist video streaming zu langsam
device = torch.device("cpu")
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(
torch.load("models/_active_model.pth", map_location='cpu'))
transformer.eval()
####################################################################################
# CAMERA CLASS (with integrated messaging queue)
####################################################################################
class Camera():
def __init__(self):
self.to_process = []
self.to_output = []
thread = threading.Thread(target=self.keep_processing, args=())
thread.daemon = True
thread.start()
def train():
parser = argparse.ArgumentParser(description='parser for style transfer')
parser.add_argument('--dataset_path',
type=str,
default=r'C:\Users\Dewey\data\celeba',
help='path to training dataset')
parser.add_argument('--style_image',
type=str,
default='mosaic.jpg',
help='path to style img')
parser.add_argument('--epochs', type=int, default=10)
parser.add_argument('--batch_size', type=int, default=4)
parser.add_argument('--image_size',
type=int,
default=256,
help='training image size')
parser.add_argument('--style_img_size',
type=int,
default=256,
help='style image size')
parser.add_argument("--lambda_content",
type=float,
default=1e5,
help="Weight for content loss")
parser.add_argument("--lambda_style",
type=float,
default=1e10,
help="Weight for style loss")
parser.add_argument('--lr', type=float, default=1e-3)
parser.add_argument("--checkpoint_model",
type=str,
help="Optional path to checkpoint model")
parser.add_argument("--checkpoint_interval",
type=int,
default=2000,
help="Batches between saving model")
parser.add_argument("--sample_interval",
type=int,
default=1000,
help="Batches between saving image samples")
parser.add_argument('--sample_format',
type=str,
default='jpg',
help='sample image format')
args = parser.parse_args()
style_name = args.style_image.split('/')[-1].split('.')[0]
os.makedirs(f'images/outputs/{style_name}-training',
exist_ok=True) # f-string格式化字符串
os.makedirs('checkpoints', exist_ok=True)
def save_sample(batch):
transformer.eval()
with torch.no_grad():
output = transformer(image_samples.to(device))
img_grid = denormalize(
torch.cat((image_samples.cpu(), output.cpu()), 2))
save_image(img_grid,
f"images/outputs/{style_name}-training/{batch}.jpg",
nrow=4)
transformer.train()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_dataset = datasets.ImageFolder(args.dataset_path,
train_transform(args.image_size))
dataloader = DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=True)
transformer = TransformerNet().to(device)
vgg = VGG16(requires_grad=False).to(device)
if args.checkpoint_model:
transformer.load_state_dict(torch.load(args.checkpoint_model))
optimizer = Adam(transformer.parameters(), lr=args.lr)
l2_loss = nn.MSELoss().to(device)
# load style image
style = style_transform(args.style_img_size)(Image.open(args.style_image))
style = style.repeat(args.batch_size, 1, 1, 1).to(device)
# style_image features
style_features = vgg(style)
gram_style = [gram(x) for x in style_features]
# visualization the image
image_samples = []
for path in random.sample(
glob.glob(f'{args.dataset_path}/*/*.{args.sample_format}'), 8):
image_samples += [
style_transform(
(args.image_size, args.image_size))(Image.open(path))
]
image_samples = torch.stack(image_samples)
c_loss = 0
s_loss = 0
t_loss = 0
for epoch in range(args.epochs):
for i, (img, _) in enumerate(dataloader):
optimizer.zero_grad()
image_original = img.to(device)
image_transformed = transformer(image_original)
origin_features = vgg(image_original)
transformed_features = vgg(image_transformed)
content_loss = args.lambda_content * l2_loss(
transformed_features.relu_2_2, origin_features.relu_2_2)
style_loss = 0
for ii, jj in zip(transformed_features, gram_style):
gram_t_features = gram(ii)
style_loss += l2_loss(gram_t_features,
jj[:img.size(), :, :]) # buyiyang
style_loss *= args.lambda_style
loss = content_loss + style_loss
loss.backward()
optimizer.step()
c_loss += content_loss.item()
s_loss += style_loss.item()
t_loss += loss.item()
print(
'[Epoch %d/%d] [Batch %d/%d] [Content: %.2f (%.2f) Style: %.2f (%.2f) Total: %.2f (%.2f)]'
% (
epoch + 1,
args.epochs,
i,
len(train_dataset),
content_loss.item(),
np.mean(c_loss),
style_loss.item(),
np.mean(s_loss),
loss.item(),
np.mean(t_loss),
))
batches_done = epoch * len(dataloader) + i + 1
if batches_done % args.sample_interval == 0:
save_sample(batches_done)
if args.checkpoint_interval > 0 and batches_done % args.checkpoint_interval == 0:
style_name = os.path.basename(args.style_image).split(".")[0]
torch.save(transformer.state_dict(),
f"checkpoints/{style_name}_{batches_done}.pth")
import torch
import torchvision
from torchvision import datasets, models, transforms
from models import TransformerNet, VGG16
import coremltools as ct
import urllib
from PIL import Image
import matplotlib.pyplot as plt
from torchvision.utils import save_image
import numpy as np
net = TransformerNet()
net.load_state_dict(
torch.load(f"starry_night_10000.pth", map_location=torch.device('cpu')))
net.eval()
x = torch.rand(1, 3, 512, 512)
traced_model = torch.jit.trace(net, x)
model = ct.convert(traced_model,
inputs=[ct.ImageType(name="input_image", shape=x.shape)])
model.save("starry_night.mlmodel")
from torchvision.utils import save_image
from PIL import Image
if __name__ == "__main__":
image_path = sys.argv[1]
model_path = sys.argv[2]
os.makedirs("images/outputs", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = style_transform()
# 定义生成网络并且加载训练好的模型
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(model_path))
transformer.eval()
# 将要进行转换的内容图片进行张量化准备
image_tensor = Variable(transform(Image.open(image_path))).to(device)
image_tensor = image_tensor.unsqueeze(0)
# 进行风格迁移
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
# 保存图片
fn = image_path.split("/")[-1]
sn = model_path.split("/")[-1]
save_image(stylized_image, f"images/outputs/{sn}-{fn}")
def main():
#Upload images
uploaded_file = st.file_uploader("Choose a picture", type=['jpg', 'png'])
# if uploaded_file is not None:
# st.image(uploaded_file, width=200)
second_uploaded_file = st.file_uploader("Choose another picture",
type=['jpg', 'png'])
# if second_uploaded_file is not None:
# st.image(second_uploaded_file, width=200)
try:
image1 = Image.open(uploaded_file)
image2 = Image.open(second_uploaded_file)
image1_arr = np.array(image1)
image2_arr = np.array(image2)
print(image1_arr.shape)
print(image2_arr.shape)
show_file = st.empty()
show_file1 = st.empty()
show_file2 = st.empty()
show_file3 = st.empty()
show_file4 = st.empty()
show_file5 = st.empty()
if not uploaded_file:
show_file.info('Please upload a file')
return
show_file.title('Input Images')
show_file1.image(uploaded_file, width=100)
show_file2.title('+')
show_file3.image(second_uploaded_file, width=100)
show_file4.title('=')
show_file5.image(image1, width=300)
# Read images to opencv
# src_img = cv2.imencode('jpg', image1)
# dst_img = cv2.imencode('jpg',image2)
except:
show_file = st.empty()
show_file.info('Please upload a file')
os.makedirs("images/outputs", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = style_transform()
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(
torch.load(
'/home/nick/Downloads/faceswap_app/models/starry_night_10000.pth'))
transformer.eval()
# Prepare input
image_tensor = Variable(transform(Image.open(uploaded_file))).to(device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
# Save image
fn = args.image_path.split("/")[-1]
save_image(stylized_image, f"images/outputs/stylized-{fn}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--image_path", type=str, required=True, help="Path to image")
parser.add_argument("--checkpoint_model", type=str, required=True, help="Path to checkpoint model")
args = parser.parse_args()
print(args)
os.makedirs("images/outputs", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = style_transform()
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(args.checkpoint_model, map_location='cpu'))
# transformer.load_state_dict(torch.load(args.checkpoint_model))
transformer.eval()
# Prepare input
image_tensor = Variable(transform(Image.open(args.image_path))).to(device)
image_tensor = image_tensor.unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = denormalize(transformer(image_tensor)).cpu()
# Save image
fn = args.image_path.split("/")[-1]
fn1 = args.checkpoint_model.split("/")[-1].split("_")[0]
save_image(stylized_image, f"./images/outputs/stylized-{fn1}-{fn}")
parser.add_argument("--checkpoint_model",
type=str,
required=True,
help="Path to checkpoint model")
args = parser.parse_args()
print(args)
os.makedirs("images/outputs", exist_ok=True)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = style_transform()
# Define model and load model checkpoint
transformer = TransformerNet().to(device)
transformer.load_state_dict(torch.load(args.checkpoint_model, 'cpu'))
transformer.eval()
stylized_frames = []
for frame in tqdm.tqdm(extract_frames(args.video_path),
desc="Processing frames"):
# Prepare input frame
image_tensor = Variable(transform(frame)).to(device).unsqueeze(0)
# Stylize image
with torch.no_grad():
stylized_image = transformer(image_tensor)
# Add to frames
stylized_frames += [deprocess(stylized_image)]
# Create video from frames
video_name = args.video_path.split("/")[-1].split(".")[0]
import torch
from utils import *
from models import TransformerNet
import os
from torchvision import transforms
import time
import cv2
from cv2 import VideoWriter, VideoWriter_fourcc
TITLE = 'gogh'
STYLE_TRANSFORM_PATH = "gogh.pth"
device = ("cuda" if torch.cuda.is_available() else "cpu")
net = TransformerNet()
net.load_state_dict(
torch.load(STYLE_TRANSFORM_PATH, map_location=torch.device(device)))
net = net.to(device)
videofile = "input.avi"
#videofile = 0
cap = cv2.VideoCapture(videofile)
if videofile is not 0:
w = int(cap.get(3))
h = int(cap.get(4))
fps = cap.get(5)
fourcc = VideoWriter_fourcc(*'mp4v')
writer = VideoWriter('output.mp4', fourcc, fps, (w, h))
count = 0
while cap.isOpened():