(load_folder, save_folder, model_path)
| 23 | |
| 24 | |
| 25 | def cartoonize(load_folder, save_folder, model_path): |
| 26 | input_photo = tf.placeholder(tf.float32, [1, None, None, 3]) |
| 27 | network_out = network.unet_generator(input_photo) |
| 28 | final_out = guided_filter.guided_filter(input_photo, network_out, r=1, eps=5e-3) |
| 29 | |
| 30 | all_vars = tf.trainable_variables() |
| 31 | gene_vars = [var for var in all_vars if 'generator' in var.name] |
| 32 | saver = tf.train.Saver(var_list=gene_vars) |
| 33 | |
| 34 | config = tf.ConfigProto() |
| 35 | config.gpu_options.allow_growth = True |
| 36 | sess = tf.Session(config=config) |
| 37 | |
| 38 | sess.run(tf.global_variables_initializer()) |
| 39 | saver.restore(sess, tf.train.latest_checkpoint(model_path)) |
| 40 | name_list = os.listdir(load_folder) |
| 41 | for name in tqdm(name_list): |
| 42 | try: |
| 43 | load_path = os.path.join(load_folder, name) |
| 44 | save_path = os.path.join(save_folder, name) |
| 45 | image = cv2.imread(load_path) |
| 46 | image = resize_crop(image) |
| 47 | batch_image = image.astype(np.float32)/127.5 - 1 |
| 48 | batch_image = np.expand_dims(batch_image, axis=0) |
| 49 | output = sess.run(final_out, feed_dict={input_photo: batch_image}) |
| 50 | output = (np.squeeze(output)+1)*127.5 |
| 51 | output = np.clip(output, 0, 255).astype(np.uint8) |
| 52 | cv2.imwrite(save_path, output) |
| 53 | except: |
| 54 | print('cartoonize {} failed'.format(load_path)) |
| 55 | |
| 56 | |
| 57 |
no test coverage detected