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)