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,
}
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('--save_intermediate', action='store_true', default=False)
parser.add_argument('--fast', action='store_true', default=False)
parser.add_argument('--no_post', action='store_true', default=False)
parser.add_argument('--cuda', type=int, default=1, help='Enable CUDA.')
args = parser.parse_args()
# Load model
p_wct = PhotoWCT()
p_wct.load_state_dict(torch.load(args.model))
if args.fast:
from photo_gif import GIFSmoothing
p_pro = GIFSmoothing(r=35, eps=0.001)
else:
from photo_smooth import Propagator
p_pro = Propagator()
if args.cuda:
p_wct.cuda(0)
process_stylization.stylization(stylization_module=p_wct,
smoothing_module=p_pro,
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,
save_intermediate=args.save_intermediate,
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):
os.makedirs(output_path)