class PhotoStyle():
def __init__(self, model_path=MODEL_PATH, use_cuda=True):
self.p_wct = PhotoWCT()
self.p_wct.load_state_dict(torch.load(model_path))
self.use_cuda = use_cuda
if use_cuda:
self.p_wct.cuda(0)
def stylize(self,
content_image_path,
style_image_path,
output_image_path,
content_seg_path=None,
style_seg_path=None,
smooth=True,
verbose=False):
process_stylization.stylization(
p_wct=self.p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_seg_path,
style_seg_path=style_seg_path,
output_image_path=output_image_path,
cuda=self.use_cuda,
smooth=smooth,
verbose=verbose,
)
def load_model(model='./PhotoWCTModels/photo_wct.pth', fast=True, cuda=1):
"""load model, fast=lighter version"""
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(model))
if fast:
from photo_gif import GIFSmoothing
p_pro = GIFSmoothing(r=35, eps=0.001)
else:
from photo_smooth import Propagator
p_pro = Propagator()
if cuda:
p_wct.cuda(0)
return p_wct, p_pro
def setup(opts):
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(PRETRAINED_MODEL_PATH))
if opts['propagation_mode'] == 'fast':
from photo_gif import GIFSmoothing
p_pro = GIFSmoothing(r=35, eps=0.001)
else:
from photo_smooth import Propagator
p_pro = Propagator()
if torch.cuda.is_available():
p_wct.cuda(0)
return {
'p_wct': p_wct,
'p_pro': p_pro,
}
class StyleTransfer_Engine:
def __init__(self):
# Load model
self.p_wct = PhotoWCT()
self.p_wct.load_state_dict(
torch.load('./PhotoWCTModels/photo_wct.pth'))
self.p_wct.cuda(0)
self.p_pro = Propagator()
def run(self, content, style):
out = stylization(stylization_module=self.p_wct,
smoothing_module=self.p_pro,
cont_img=content,
styl_img=style,
cuda=1,
save_intermediate=False,
no_post=False)
return out
def main():
parser = argparse.ArgumentParser(
description='Photorealistic Image Stylization')
parser.add_argument(
'--model',
default='./PhotoWCTModels/photo_wct.pth',
help=
'Path to the PhotoWCT model. These are provided by the PhotoWCT submodule, please use `git submodule update --init --recursive` to pull.'
)
parser.add_argument('--content_image_path',
default='./images/content1.png')
parser.add_argument('--content_seg_path', default=[])
parser.add_argument('--style_image_path', default='./images/style1.png')
parser.add_argument('--style_seg_path', default=[])
parser.add_argument('--output_image_path',
default='./results/example1.png')
parser.add_argument('--cuda', type=int, default=1, help='Enable CUDA.')
args = parser.parse_args()
# Load model
p_wct = PhotoWCT()
try:
p_wct.load_state_dict(torch.load(args.model))
except:
print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
exit()
if args.cuda:
p_wct.cuda(0)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=args.content_image_path,
style_image_path=args.style_image_path,
content_seg_path=args.content_seg_path,
style_seg_path=args.style_seg_path,
output_image_path=args.output_image_path,
cuda=args.cuda,
)
folder = 'examples'
cont_img_folder = os.path.join(folder, 'content_img')
cont_seg_folder = os.path.join(folder, 'content_seg')
styl_img_folder = os.path.join(folder, 'style_img')
styl_seg_folder = os.path.join(folder, 'style_seg')
outp_img_folder = os.path.join(folder, 'results')
cont_img_list = [
f for f in os.listdir(cont_img_folder)
if os.path.isfile(os.path.join(cont_img_folder, f))
]
cont_img_list.sort()
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model))
p_wct.cuda(0)
for f in cont_img_list:
print("Process " + f)
content_image_path = os.path.join(cont_img_folder, f)
content_seg_path = os.path.join(cont_seg_folder, f).replace(".png", ".pgm")
style_image_path = os.path.join(styl_img_folder, f)
style_seg_path = os.path.join(styl_seg_folder, f).replace(".png", ".pgm")
output_image_path = os.path.join(outp_img_folder, f)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
from flask import Flask, request
from flask import jsonify
import process_stylization
from image_utils import base64_png_image_to_pillow_image, get_apt_image_size, get_temp_png_file_path
from photo_wct import PhotoWCT
app = Flask(__name__)
MAX_NUM_PIXELS = 1024 * 512
# Load model
p_wct = PhotoWCT()
try:
p_wct.load_state_dict(
torch.load(
os.path.join(os.path.dirname(__file__), 'PhotoWCTModels',
'photo_wct.pth')))
except:
print("Fail to load PhotoWCT models. PhotoWCT submodule not updated?")
exit()
@app.route("/stylize/", methods=['POST'])
def stylize():
content_image_base64 = request.json.get('content_image_base64', None)
if content_image_base64 is None:
raise Exception('content_image_base64 cannot be None')
else:
content_image = base64_png_image_to_pillow_image(content_image_base64)
style_image_base64 = request.json.get('style_image_base64', None)
help='Path to the PhotoWCT model. These are provided by the PhotoWCT submodule, please use `git submodule update --init --recursive` to pull.')
parser.add_argument('--cuda', type=bool, default=True, help='Enable CUDA.')
args = parser.parse_args()
folder = 'examples'
cont_img_folder = os.path.join(folder, 'content_img')
cont_seg_folder = os.path.join(folder, 'content_seg')
styl_img_folder = os.path.join(folder, 'style_img')
styl_seg_folder = os.path.join(folder, 'style_seg')
outp_img_folder = os.path.join(folder, 'results')
cont_img_list = [f for f in os.listdir(cont_img_folder) if os.path.isfile(os.path.join(cont_img_folder, f))]
cont_img_list.sort()
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model))
for f in cont_img_list:
print("Process " + f)
content_image_path = os.path.join(cont_img_folder, f)
content_seg_path = os.path.join(cont_seg_folder, f).replace(".png", ".pgm")
style_image_path = os.path.join(styl_img_folder, f)
style_seg_path = os.path.join(styl_seg_folder, f).replace(".png", ".pgm")
output_image_path = os.path.join(outp_img_folder, f)
process_stylization.stylization(
p_wct=p_wct,
content_image_path=content_image_path,
style_image_path=style_image_path,
content_seg_path=content_seg_path,
from __future__ import print_function
from model import StyleTransferModel
from io import BytesIO
from telegram import ReplyKeyboardMarkup, ReplyKeyboardRemove
from telegram.ext.dispatcher import run_async
import torch
from PIL import Image
import process_stylization
from photo_wct import PhotoWCT
from photo_gif import GIFSmoothing
model = StyleTransferModel()
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load('photo_wct.pth'))
p_pro = GIFSmoothing(r=35, eps=0.001)
first_image_file = {}
storage = []
storage1 = []
@run_async
def send_prediction_on_photo(update, context):
global storage
global storage1
# Нам нужно получить две картинки, чтобы произвести перенос стиля, но каждая картинка приходит в
# отдельном апдейте, поэтому в простейшем случае мы будем сохранять id первой картинки в память,
# чтобы, когда уже придет вторая, мы могли загрузить в память уже сами картинки и обработать их.
# Точно место для улучшения, я бы
bot = context.bot
if update.message.text == '/style':
storage.append('1')
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
root_dir = 'dataset/'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
train_data = datasets.ImageFolder(root_dir, image_transform)
data_loader = torch.utils.data.DataLoader(train_data,
batch_size=32,
shuffle=True)
# define & initialize model
model = PhotoWCT()
model.load_state_dict(torch.load('./PhotoWCTModels/photo_wct.pth'))
# set encoders requires_grad=False, i.e. freeze them during training
set_parameter_requires_grad(model.e1, False)
set_parameter_requires_grad(model.e2, False)
set_parameter_requires_grad(model.e3, False)
set_parameter_requires_grad(model.e4, False)
# transfer model to GPU if you have one
model.to(device)
# set criterion to reconstruction loss & define optimizer
criterion = WeightedMseContentLoss(content_loss_weight=1, mse_loss_weight=1700)
#criterion = ContentLoss() # MSELoss() # alternatives to the above dissimilarity constraint
lr = 0.0001
from os import path as osp
SEED = 1984
N_IMGS = 100
FAST = True
SPLIT = "val"
MODEL = "./PhotoWCTModels/photo_wct.pth"
STYLE_PATH = "/data/datasets/style_transfer_amos/styles_sub_10"
CONTENT_PATH = "/data/datasets/phototourism"
OUTPUT_ST_PATH = osp.join(CONTENT_PATH, "style_transfer_all")
STYLES = ["cloudy", "dusk", "mist", "night", "rainy", "snow"]
#STYLES = ["snow"]
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(MODEL))
p_pro = GIFSmoothing(r=35, eps=0.001) if FAST else Propagator()
p_wct.cuda(0)
with open(osp.join(CONTENT_PATH, SPLIT + "_phototourism_ms.txt"), "r") as f:
content_fnames = [line.rstrip('\n') for line in f]
for style in STYLES:
print("Style: {:s}".format(style))
style_fnames = [img for img in os.listdir(osp.join(STYLE_PATH, style)) if img[-3:] in ["png", "jpg"]]
for style_fname in style_fnames:
k_cont = 0
for content_fname in content_fnames:
scene = content_fname.split('/')[0]
output_path = osp.join(CONTENT_PATH, "style_transfer_all", scene, style)
if not osp.isdir(output_path):
