def transform_forward(path_in, path_out, checkpoint_dir): image_input = get_image(path_in) image_input = np.array([image_input]) image_shape = image_input.shape config = tf.ConfigProto(allow_soft_placement=True) config.gpu_options.allow_growth = True with tf.Graph().as_default(), tf.device('/cpu:0'), tf.Session(config=config) as sess: initial_image = tf.placeholder(tf.float32, shape=image_shape, name='initial_image') with tf.variable_scope("Transform_net"): transformed_image_tensor = transform.net(initial_image) saver = tf.train.Saver() if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) transformed_image = sess.run(transformed_image_tensor, feed_dict={initial_image: image_input}) save_img(path_out, transformed_image[0])
def optimize(content_targets, style_target, content_weight, style_weight, tv_weight, vgg_path, epochs=2, print_iterations=1000, batch_size=4, save_path='saver/fns.ckpt', slow=False, learning_rate=1e-3, debug=False): ### mod = len(content_targets) % batch_size # 由于batch_size=4,如果content一共有7个样本,那么只能构成1个batch if mod > 0: # 所以后面的3个就要被丢弃,这段代码就完成这样的功能! content_targets = content_targets[:-mod] # *** 这里的[:-mod]表示取content从第1个开始到倒数第mod个样本! *** ### style_features = {} batch_shape = (batch_size, 256, 256, 3) # batch_size * H * W * C;这是tensorflow的标准数据格式 style_shape = (1,) + style_target.shape # 前面的1表示style的batch_size = 1 with tf.Graph.as_default(),tf.device('/gpu:0'),tf.Session as sess: # 分别是指定:tf的计算图,运行设备,Session style_image = tf.placeholder(tf.float32, shape=style_shape, name='style_image') style_image_pre = vgg.preprocess(style_image) net = vgg.net(vgg_path, style_image_pre) # 调用来让预处理完成的style图片跑一遍VGG stylp_pre = np.array(style_target) for layer in STYLE_LAYERS: features = net[layer].eval(feed_dict = {style_image:stylp_pre}) # 取出对应层的feature_map features = np.reshape(features, (-1, features.shape[3])) ### gram = np.matmul(features.T, features) # 这里的3行代码是根据论文的设计实现而已 style_features[layer] = gram ### with tf.Graph.as_default(),tf.Session as sess: # 分别是指定:tf的计算图,运行设备,Session x_content = tf.placeholder(tf.float32, shape=batch_shape,name='x_content') x_pre = vgg.preprocess(x_content) content_features = {} content_net = vgg.net(vgg_path, x_pre) content_features[CONTENT_LAYER] = content_net[CONTENT_LAYER] preds = transform.net(x_content/255.0)
def __init__(self, vgg_path, style_image, content_shape, content_weight, style_weight, tv_weight, batch_size, device, log_f): with tf.device(device): self.log_file = log_f vgg = vgg_network.VGG(vgg_path) self.style_image = style_image self.batch_size = batch_size self.batch_shape = (batch_size, ) + content_shape self.input_batch = tf.placeholder(tf.float32, shape=self.batch_shape, name="input_batch") self.stylized_image = transform.net(self.input_batch, _vgg=vgg) loss_calculator = LossCalculator(vgg, self.stylized_image) self.content_loss = loss_calculator.content_loss( self.input_batch, self.CONTENT_LAYER, content_weight) / self.batch_size self.style_loss = loss_calculator.style_loss( self.style_image, self.STYLE_LAYERS, style_weight) / self.batch_size self.total_variation_loss = loss_calculator.tv_loss( self.stylized_image, tv_weight) / batch_size self.loss = self.content_loss + self.style_loss + self.total_variation_loss
def protobuf_from_checkpoint(ckpt_file, image_shape, batch_size, output_name): if not os.path.isfile(ckpt_file): raise ValueError(f'File "{ckpt_file}" does not exist or is not a file.') # create the tf Session sess = tf.Session() # Compute the shape of the input placeholder. This shape is what the serialized model can # process. For other input shapes you will have to resize the images or make a new model export batch_shape = [batch_size] + image_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') # create the network to the variables are in the global scope preds = transform.net(img_placeholder) # noqa saver = tf.train.Saver() # load our checkpoint into the variables saver.restore(sess, ckpt_file) # get the tf graph and retrieve operation names graph = tf.get_default_graph() op_names = [op.name for op in graph.get_operations()] # convert the protobuf GraphDef to a GraphDef that has no variables but just constants with the # current values. output_graph_def = graph_util.convert_variables_to_constants( sess, graph.as_graph_def(), op_names) # dump GraphDef to file graph_io.write_graph(output_graph_def, './', output_name, as_text=False) sess.close()
def ffwd(data_in, paths_out, checkpoint_dir, device_t='/cpu:0', bw=False): img_shape = get_img(data_in).shape g = tf.Graph() with g.as_default(), g.device(device_t), tf.Session(config=soft_config) as sess: batch_shape = (1,) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) X = np.zeros(batch_shape, dtype=np.float32) if bw: img = get_bw_img(data_in) else: img = get_img(data_in) X[0] = img _preds = sess.run(preds, feed_dict={img_placeholder:X}) save_img(paths_out, _preds[0])
def rundeeplearning(img_in, checkpoint_dir, batch_size=1): img = get_rgb_np(img_in) img_shape = get_img(img).shape g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) config_mu = tf.ConfigProto(intra_op_parallelism_threads=16) config_ = tf.ConfigProto(intra_op_parallelism_threads=1, inter_op_parallelism_threads=1, \ allow_soft_placement=True, device_count = {'CPU': 1}) soft_config.gpu_options.allow_growth = True with g.as_default(), \ tf.Session(config=config_mu) as sess: batch_shape = (batch_size, ) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() # Load saver.restore(sess, checkpoint_dir) X = np.zeros(batch_shape, dtype=np.float32) X[0] = img _preds = sess.run(preds, feed_dict={img_placeholder: X}) img_np = np.clip(_preds[0], 0, 255).astype(np.uint8) img_out = Image.fromarray(img_np) buffer = io.BytesIO() img_out.save(buffer, format="PNG") return buffer
def ckpt2pb(ckpt_path): graph = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True sess = tf.Session(config=soft_config) # construct the model first batch_shape = (1, 256, 256, 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') # transform.net is the func which build and return the model in this case # replace it use your model -> output = your_model(img_placeholder) output = transform.net(img_placeholder) graph_def = sess.graph.as_graph_def() for node in graph_def.node: print _node_name(node.name) # find the ouput tensor name of the model in the print above # assumpt it is the last one final_tensor_name =_node_name(graph_def.node[-1].name) saver = tf.train.Saver() saver.restore(sess, ckpt_path) output_filename = ckpt_path.replace('ckpt', 'pb') output_graph_def = graph_util.convert_variables_to_constants(sess, sess.graph.as_graph_def(), [final_tensor_name]) with gfile.FastGFile(output_filename, 'wb') as f: f.write(output_graph_def.SerializeToString())
def simple_evaluate(image_array, checkpoint_dir, device_t='/cpu:0'): img_shape = image_array.shape reshaped_image = image_array.reshape((1, img_shape[0], img_shape[1], img_shape[2])) g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), tf.Session(config=soft_config) as sess: batch_shape = (1,) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") _preds = sess.run(preds, feed_dict={img_placeholder: reshaped_image}) result = np.clip(_preds[0], 0, 255).astype(np.uint8) # Check if shapes match and act accordingly if result.shape == img_shape: return result else: return result[0: img_shape[0], 0: img_shape[1], 0: img_shape[2]]
def get_ckpt_model(): models_dir = os.listdir("temp") for i, style_name in enumerate(models_dir): style_name, _ = os.path.splitext(style_name) batch_shape = (None, None, None, 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) with tf.Session() as sess: saver = tf.train.Saver() print("process {}, {}/{}".format(style_name, i, len(models_dir))) old_model_dir = "temp/" + style_name + ".ckpt" new_model_dir = "./temp2/" + style_name + "/" if os.path.isdir(old_model_dir): ckpt = tf.train.get_checkpoint_state(old_model_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, old_model_dir) print(saver.save(sess, new_model_dir)) tf.compat.v1.reset_default_graph()
def main(): batch_size = 1 checkpoint_dir = 'style/wave.ckpt' device_t = '/gpu:0' data_in = 'in/chicago.jpg' path_out = 'out/chicago.jpg' img_shape = get_img(data_in).shape is_paths = type(data_in[0]) == str g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), tf.Session( config=soft_config) as sess: batch_shape = (batch_size, ) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) X = np.zeros(batch_shape, dtype=np.float32) img = get_img(data_in) X[0] = img _preds = sess.run(preds, feed_dict={img_placeholder: X}) save_img(path_out, _preds[0])
def style_transfer(style, in_path, out_path, device_t='/gpu:0', batch_size=1): print("Log: Begin - style transfer") curr_style = STYLE_MODEL_MAP[style] img_shape = get_img(in_path).shape graph = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with graph.as_default(), graph.device(device_t), tf.Session(config=soft_config) as sess: img_val = tf.placeholder(tf.float32, shape=(1,) + img_shape, name='img_val') preds = transform.net(img_val) saver = tf.train.Saver() if os.path.isdir(curr_style): ckpt = tf.train.get_checkpoint_state(curr_style) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, curr_style) img = get_img(in_path) _preds = sess.run(preds, feed_dict={img_val:[img]}) save_img(out_path, _preds[0]) print("Log: End - style transfer - Success√") return _preds[0]
def ffwd_video(path_in, path_out, checkpoint_dir, device_t='/gpu:0', batch_size=4): video_clip = VideoFileClip(path_in, audio=False) video_writer = ffmpeg_writer.FFMPEG_VideoWriter(path_out, video_clip.size, video_clip.fps, codec="libx264", preset="medium", bitrate="2000k", audiofile=path_in, threads=None, ffmpeg_params=None) g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), \ tf.Session(config=soft_config) as sess: batch_shape = (batch_size, video_clip.size[1], video_clip.size[0], 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) X = np.zeros(batch_shape, dtype=np.float32) def style_and_write(count): for i in range(count, batch_size): X[i] = X[count - 1] # Use last frame to fill X _preds = sess.run(preds, feed_dict={img_placeholder: X}) for i in range(0, count): video_writer.write_frame( np.clip(_preds[i], 0, 255).astype(np.uint8)) frame_count = 0 # The frame count that written to X for frame in video_clip.iter_frames(): X[frame_count] = frame frame_count += 1 if frame_count == batch_size: style_and_write(frame_count) frame_count = 0 if frame_count != 0: style_and_write(frame_count) video_writer.close()
def ffwd(data_in, paths_out, checkpoint_dir, device_t='/gpu:0', batch_size=4): assert len(paths_out) > 0 is_paths = type(data_in[0]) == str if is_paths: assert len(data_in) == len(paths_out) img_shape = get_img(data_in[0]).shape else: assert data_in.size[0] == len(paths_out) img_shape = X[0].shape g = tf.Graph() batch_size = min(len(paths_out), batch_size) curr_num = 0 soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), \ tf.Session(config=soft_config) as sess: batch_shape = (batch_size,) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) num_iters = int(len(paths_out)/batch_size) for i in range(num_iters): pos = i * batch_size curr_batch_out = paths_out[pos:pos+batch_size] if is_paths: curr_batch_in = data_in[pos:pos+batch_size] X = np.zeros(batch_shape, dtype=np.float32) for j, path_in in enumerate(curr_batch_in): img = get_img(path_in) assert img.shape == img_shape, \ 'Images have different dimensions. ' + \ 'Resize images or use --allow-different-dimensions.' X[j] = img else: X = data_in[pos:pos+batch_size] _preds = sess.run(preds, feed_dict={img_placeholder:X}) for j, path_out in enumerate(curr_batch_out): save_img(path_out, _preds[j]) remaining_in = data_in[num_iters*batch_size:] remaining_out = paths_out[num_iters*batch_size:] if len(remaining_in) > 0: ffwd(remaining_in, remaining_out, checkpoint_dir, device_t=device_t, batch_size=1)
def deep_art(img, style): # check if style is valid if not style in [ "la_muse", "rain_princess", "scream", "udnie", "wave", "wreck" ]: return "Please enter a valid style from the list: [la_muse, rain_princess, scream, udnie, wave, wreck]" # whether to use GPU or CPU, setup batch size, and select which style to use device_t = '/gpu:0' batch_size = 1 checkpoint_dir = f"./fast-style-transfer/models/{style}.ckpt" img_shape = img.shape g = tf.Graph() curr_num = 0 # setup GPU configurations for model soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True # setup batch shape from image size to be compatible with model input shape with g.as_default(), g.device(device_t), \ tf.Session(config=soft_config) as sess: batch_shape = (batch_size, ) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') # 'predict' the output stylized image and save it pred = transform.net(img_placeholder) saver = tf.train.Saver() # use pretrained model as the checkpoint to start from if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) # in this case we do 1 image at a time, so num_iters = 1 num_iters = 1 pos = batch_size # change image dimensions to 1xNxMx3 from NxMx3 then run the model styled = sess.run(pred, feed_dict={img_placeholder: img[np.newaxis, ...]}) # change output back to 1xNxMx3 and convert to image im = Image.fromarray(styled[0, :, :].astype('uint8'), 'RGB') return im
def evaluate_img(img_in, img_path, ckpt): img_shape = (256, 256, 3) batch_shape = (1, 256, 256, 3) soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with tf.Graph().as_default(), tf.Session(config=soft_config) as sess: # Declare placeholders we'll feed into the graph X_inputs = tf.placeholder( tf.float32, shape=batch_shape, name='X_inputs') # Define output node preds = transform.net(X_inputs) # (1, 720, 720, 3) tf.identity(preds[0], name='output') # For restore training checkpoints (important) saver = tf.train.Saver() saver.restore(sess, ckpt) # run """ saver = tf.train.Saver() if os.path.isdir(FLAGS.checkpoint_dir): ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) # run else: raise Exception("No checkpoint found...") else: ckpt = saver.restore(sess, FLAGS.checkpoint_dir) """ X = np.zeros(batch_shape, dtype=np.float32) # feed img = get_img(img_in, img_shape) X[0] = img _preds = sess.run(preds, feed_dict={X_inputs: X}) save_img(img_path, _preds[0]) # Write graph. start_time = time.time() tf.train.write_graph( sess.graph.as_graph_def(), FLAGS.model_dir, FLAGS.model_name + '.pb', as_text=False) tf.train.write_graph( sess.graph.as_graph_def(), FLAGS.model_dir, FLAGS.model_name + '.pb.txt', as_text=True) end_time = time.time() delta_time = end_time - start_time print('Save pb and pb.txt done!, time:', delta_time)
def ffwd(content, network_path): with tf.Session() as sess: content_placeholder = tf.placeholder(tf.float32, shape=content.shape, name='content_placeholder') network = transform.net(content_placeholder / 255.0) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(network_path) saver.restore(sess, ckpt.model_checkpoint_path) prediction = sess.run(network, feed_dict={content_placeholder: content}) return prediction[0]
def __init__(self, image_resolution=(300, 250)): self._image_resolution = image_resolution self._graph = tf.Graph() self._session = tf.Session(graph=self._graph) with self._graph.as_default(): self._image_placeholder = tf.placeholder(dtype=tf.float32, shape=(1, 300, 250, 3), name='img_placeholder') self._output = transform.net(self._image_placeholder)
def generate(): if not FLAGS.CONTENT_IMAGE: tf.logging.info("train a fast nerual style need to set the Content images path") return if not os.path.exists(FLAGS.OUTPUT_FOLDER): os.mkdir(FLAGS.OUTPUT_FOLDER) # 获取图片信息 height = 0 width = 0 with open(FLAGS.CONTENT_IMAGE, 'rb') as img: with tf.Session().as_default() as sess: if FLAGS.CONTENT_IMAGE.lower().endswith('png'): image = sess.run(tf.image.decode_png(img.read())) else: image = sess.run(tf.image.decode_jpeg(img.read())) height = image.shape[0] width = image.shape[1] tf.logging.info('Image size: %dx%d' % (width, height)) with tf.Graph().as_default(), tf.Session() as sess: # 读取图片文件 path = FLAGS.CONTENT_IMAGE png = path.lower().endswith('png') img_bytes = tf.read_file(path) # 图片解码 content_image = tf.image.decode_png(img_bytes, channels=3) if png else tf.image.decode_jpeg(img_bytes, channels=3) content_image = tf.image.convert_image_dtype(content_image, tf.float32) * 255.0 content_image = tf.expand_dims(content_image, 0) generated_images = transform.net(content_image - vgg.MEAN_PIXEL, training=False) output_format = tf.saturate_cast(generated_images, tf.uint8) # 开始转换 saver = tf.train.Saver(tf.global_variables(), write_version=tf.train.SaverDef.V1) sess.run([tf.global_variables_initializer(), tf.local_variables_initializer()]) model_path = os.path.abspath(FLAGS.MODEL_PATH) tf.logging.info('Usage model {}'.format(model_path)) saver.restore(sess, model_path) filename = os.path.basename(FLAGS.CONTENT_IMAGE) (shotname, extension) = os.path.splitext(filename) filename = shotname + '-' + os.path.basename(FLAGS.MODEL_PATH) + extension tf.logging.info("image {}".format(filename)) images_t = sess.run(output_format) assert len(images_t) == 1 misc.imsave(os.path.join(FLAGS.OUTPUT_FOLDER, filename), images_t[0])
def __init__(self, checkpoint_dir, device_t): # Create local graph and use it in the session self.graph = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True self.sess = tf.Session(graph=self.graph, config=soft_config) with self.graph.as_default(), self.graph.device(device_t): batch_shape = (1,) + (256, 256, 3) self.img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') self.pred = transform.net(self.img_placeholder) # Import saved model from location 'checkpoint_dir' into local graph saver = tf.train.Saver() saver.restore(self.sess, checkpoint_dir)
def setup(options): global sess global img_placeholder global preds global g h, w = 480, 640 img_shape = (h, w, 3) batch_shape = (1,) + img_shape g = tf.get_default_graph() sess = tf.Session(graph=g) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) load_checkpoint(os.path.join(options['checkpoint_path'], 'fns.ckpt'), sess) return sess
def forward_prop(data_in, paths_out, checkpoint_dir, device_t='/cpu:0', batch_size=1): assert len(paths_out) > 0 is_paths = type(data_in) == str img_shape = get_img(data_in).shape #print("Batch size: ", batch_size) g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), tf.Session( config=soft_config) as sess: batch_shape = (batch_size, ) + img_shape #print("Batch_shape: ", batch_shape) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() #Restore checkpoint in session saver.restore(sess, checkpoint_dir) curr_batch_out = paths_out if is_paths: curr_batch_in = data_in print("curr_batch_in: ", curr_batch_in) print("curr_batch_out: ", curr_batch_out) X = np.zeros(batch_shape, dtype=np.float32) img = get_img(curr_batch_in) assert img.shape == img_shape, 'Images have different dimensions. ' + 'Resize images' X[0] = img #print("Shape: ", X.shape) #(1,960,960,3) _preds = sess.run(preds, feed_dict={img_placeholder: X}) save_img(curr_batch_out, _preds[0]) sess.close() print("Done!!") return curr_batch_out
def main(checkpoint_path, input_shape, out_graph_name): # Init graph and session to be used g = tf.Graph() soft_config = tf.compat.v1.ConfigProto(allow_soft_placement=True) with g.as_default(), g.device('/cpu'), tf.compat.v1.Session( config=soft_config) as sess: # Placeholder variable for graph input img_placeholder = tf.compat.v1.placeholder(tf.float32, shape=input_shape, name='img_placeholder') # The model from the repo transform.net(img_placeholder) # Restore model from checkpoint saver = tf.compat.v1.train.Saver() saver.restore(sess, checkpoint_path) # Freeze graph from the session. # "add_37" is the actual last operation of graph frozen = tf.compat.v1.graph_util.convert_variables_to_constants( sess, sess.graph_def, ["add_37"]) # Write frozen graph to a file graph_io.write_graph(frozen, './', out_graph_name, as_text=False) print(f'Frozen graph {out_graph_name} is saved!')
def ffwd_video(path_in, path_out, checkpoint_dir, device_t='/gpu:0', batch_size=4): video_clip = VideoFileClip(path_in, audio=False) video_writer = ffmpeg_writer.FFMPEG_VideoWriter(path_out, video_clip.size, video_clip.fps, codec="libx264", preset="medium", bitrate="2000k", audiofile=path_in, threads=None, ffmpeg_params=None) g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), \ tf.Session(config=soft_config) as sess: batch_shape = (batch_size, video_clip.size[1], video_clip.size[0], 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() if os.path.isdir(checkpoint_dir): ckpt = tf.train.get_checkpoint_state(checkpoint_dir) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, checkpoint_dir) X = np.zeros(batch_shape, dtype=np.float32) def style_and_write(count): for i in range(count, batch_size): X[i] = X[count - 1] # Use last frame to fill X _preds = sess.run(preds, feed_dict={img_placeholder: X}) for i in range(0, count): video_writer.write_frame(np.clip(_preds[i], 0, 255).astype(np.uint8)) frame_count = 0 # The frame count that written to X for frame in video_clip.iter_frames(): X[frame_count] = frame frame_count += 1 if frame_count == batch_size: style_and_write(frame_count) frame_count = 0 if frame_count != 0: style_and_write(frame_count) video_writer.close()
def ffwd(content, network_path, cin): with tf.Session() as sess: init = tf.global_variables_initializer() sess.run(init) img_placeholder = tf.placeholder(tf.float32, shape=content.shape, name='img_placeholder') network = transform.net(img_placeholder) saver = tf.train.Saver() saver.restore(sess, tf.train.latest_checkpoint(network_path)) prediction = sess.run(network, feed_dict={img_placeholder: content}) sess.close() return prediction[0]
def setup(options): global sess global img_placeholder global preds h, w = 480, 640 img_shape = (h, w, 3) batch_shape = (1,) + img_shape g = tf.Graph() g.as_default() g.device('/gpu:0') soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True sess = tf.Session(config=soft_config) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) load_checkpoint(options['checkpoint_path'], sess) return sess
def create_transfer(img, checkpoint_dir, device_t='/gpu:0'): g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), \ tf.Session(config=soft_config) as sess: batch_shape = (1,) + img.shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(checkpoint_dir) saver.restore(sess, ckpt.model_checkpoint_path) X = np.zeros(batch_shape, dtype=np.float32) X[0] = img _preds = sess.run(preds, feed_dict={img_placeholder:X}) return _preds
def ffwd(content, network_path): with tf.Session() as sess: img_placeholder = tf.placeholder(tf.float32, shape=content.shape, name='img_placeholder') network = transform.net(img_placeholder) saver = tf.train.Saver() ckpt = tf.train.get_checkpoint_state(network_path) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") prediction = sess.run(network, feed_dict={img_placeholder: content}) return prediction[0]
def convert(model_path, input_paths, output_paths): sess = tf.Session(config=config) batch_shape = (None, None, None, 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() saver.restore(sess, model_path) for input_path, output_path in zip(input_paths, output_paths): # Transform image x_input = get_img(input_path) x_input = np.expand_dims(x_input, 0) y_out = sess.run(preds, feed_dict={img_placeholder: x_input}) save_img(output_path, y_out[0]) tf.compat.v1.reset_default_graph() gc.collect() return
def ffwd(content, mask, model): with tf.Session() as sess: shape = list(content.shape) img_placeholder = tf.placeholder(tf.float32, shape=shape, name='img_placeholder') print(img_placeholder.shape) mask_placeholder = tf.placeholder(tf.float32, shape=[1, 64, 64, 1]) network = transform.net(img_placeholder, mask_placeholder) saver = tf.train.Saver() saver.restore(sess, model) prediction = sess.run(network, feed_dict={ img_placeholder: content, mask_placeholder: mask }) return prediction[0]
def export(checkpoint_dir, batch_shape): with tf.Graph().as_default(), tf.Session() as sess: images = tf.placeholder(tf.float32, shape=batch_shape, name='images') preds = transform.net(images / 255.0) tensor_info_images = tf.saved_model.utils.build_tensor_info(images) tensor_info_preds = tf.saved_model.utils.build_tensor_info(preds) prediction_signature = ( tf.saved_model.signature_def_utils.build_signature_def( inputs={'images': tensor_info_images}, outputs={'preds': tensor_info_preds}, method_name=tf.saved_model.signature_constants. PREDICT_METHOD_NAME)) saver = tf.train.Saver() saver.restore(sess, tf.train.latest_checkpoint(checkpoint_dir)) export_path = os.path.join(checkpoint_dir, "1") builder = tf.saved_model.builder.SavedModelBuilder(export_path) builder.add_meta_graph_and_variables( sess, [tf.saved_model.tag_constants.SERVING], signature_def_map={'predict_images': prediction_signature}) builder.save()
def ffwd(checkpoint_dir, device_t='/gpu:0'): img_shape = (500, 500, 3) g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(device_t), \ tf.Session(config=soft_config) as sess: batch_shape = (1, ) + img_shape img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() saver.restore(sess, checkpoint_dir) import cv2 video_stream = cv2.VideoCapture(0) while (True): if cv2.waitKey(1) & 0xFF == ord('q'): break ret, frame = video_stream.read() img = cv2.resize(frame, (img_shape[0], img_shape[1])) cv2.imshow('frame', img) #img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) X = [img] _preds = sess.run(preds, feed_dict={img_placeholder: X}) # save_img(path_out, _preds[j #print(_preds[0].shape) result = np.clip(_preds[0], 0, 255).astype(np.uint8) cv2.imshow('frame', cv2.resize(result, (img_shape[0], img_shape[1])))
def from_pipe(opts): command = ["ffprobe", '-v', "quiet", '-print_format', 'json', '-show_streams', opts.in_path] info = json.loads(str(subprocess.check_output(command), encoding="utf8")) width = int(info["streams"][0]["width"]) height = int(info["streams"][0]["height"]) fps = round(eval(info["streams"][0]["r_frame_rate"])) command = ["ffmpeg", '-loglevel', "quiet", '-i', opts.in_path, '-f', 'image2pipe', '-pix_fmt', 'rgb24', '-vcodec', 'rawvideo', '-'] pipe_in = subprocess.Popen(command, stdout=subprocess.PIPE, bufsize=10 ** 9, stdin=None, stderr=None) command = ["ffmpeg", '-loglevel', "info", '-y', # (optional) overwrite output file if it exists '-f', 'rawvideo', '-vcodec', 'rawvideo', '-s', str(width) + 'x' + str(height), # size of one frame '-pix_fmt', 'rgb24', '-r', str(fps), # frames per second '-i', '-', # The imput comes from a pipe '-an', # Tells FFMPEG not to expect any audio '-c:v', 'libx264', '-preset', 'slow', '-crf', '18', opts.out] pipe_out = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=None, stderr=None) g = tf.Graph() soft_config = tf.ConfigProto(allow_soft_placement=True) soft_config.gpu_options.allow_growth = True with g.as_default(), g.device(opts.device), \ tf.Session(config=soft_config) as sess: batch_shape = (opts.batch_size, height, width, 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) saver = tf.train.Saver() if os.path.isdir(opts.checkpoint): ckpt = tf.train.get_checkpoint_state(opts.checkpoint) if ckpt and ckpt.model_checkpoint_path: saver.restore(sess, ckpt.model_checkpoint_path) else: raise Exception("No checkpoint found...") else: saver.restore(sess, opts.checkpoint) X = np.zeros(batch_shape, dtype=np.float32) nbytes = 3 * width * height read_input = True last = False while read_input: count = 0 while count < opts.batch_size: raw_image = pipe_in.stdout.read(width * height * 3) if len(raw_image) != nbytes: if count == 0: read_input = False else: last = True X = X[:count] batch_shape = (count, height, width, 3) img_placeholder = tf.placeholder(tf.float32, shape=batch_shape, name='img_placeholder') preds = transform.net(img_placeholder) break image = numpy.fromstring(raw_image, dtype='uint8') image = image.reshape((height, width, 3)) X[count] = image count += 1 if read_input: if last: read_input = False _preds = sess.run(preds, feed_dict={img_placeholder: X}) for i in range(0, batch_shape[0]): img = np.clip(_preds[i], 0, 255).astype(np.uint8) try: pipe_out.stdin.write(img) except IOError as err: ffmpeg_error = pipe_out.stderr.read() error = (str(err) + ("\n\nFFMPEG encountered" "the following error while writing file:" "\n\n %s" % ffmpeg_error)) read_input = False print(error) pipe_out.terminate() pipe_in.terminate() pipe_out.stdin.close() pipe_in.stdout.close() del pipe_in del pipe_out
def optimize(content_targets, style_target, content_weight, style_weight, tv_weight, vgg_path, epochs=2, print_iterations=1000, batch_size=4, save_path='saver/fns.ckpt', slow=False, learning_rate=1e-3, debug=False): if slow: batch_size = 1 mod = len(content_targets) % batch_size if mod > 0: print("Train set has been trimmed slightly..") content_targets = content_targets[:-mod] style_features = {} batch_shape = (batch_size,256,256,3) style_shape = (1,) + style_target.shape print(style_shape) # precompute style features with tf.Graph().as_default(), tf.device('/cpu:0'), tf.Session() as sess: style_image = tf.placeholder(tf.float32, shape=style_shape, name='style_image') style_image_pre = vgg.preprocess(style_image) net = vgg.net(vgg_path, style_image_pre) style_pre = np.array([style_target]) for layer in STYLE_LAYERS: features = net[layer].eval(feed_dict={style_image:style_pre}) features = np.reshape(features, (-1, features.shape[3])) gram = np.matmul(features.T, features) / features.size style_features[layer] = gram with tf.Graph().as_default(), tf.Session() as sess: X_content = tf.placeholder(tf.float32, shape=batch_shape, name="X_content") X_pre = vgg.preprocess(X_content) # precompute content features content_features = {} content_net = vgg.net(vgg_path, X_pre) content_features[CONTENT_LAYER] = content_net[CONTENT_LAYER] if slow: preds = tf.Variable( tf.random_normal(X_content.get_shape()) * 0.256 ) preds_pre = preds else: preds = transform.net(X_content/255.0) preds_pre = vgg.preprocess(preds) net = vgg.net(vgg_path, preds_pre) content_size = _tensor_size(content_features[CONTENT_LAYER])*batch_size assert _tensor_size(content_features[CONTENT_LAYER]) == _tensor_size(net[CONTENT_LAYER]) content_loss = content_weight * (2 * tf.nn.l2_loss( net[CONTENT_LAYER] - content_features[CONTENT_LAYER]) / content_size ) style_losses = [] for style_layer in STYLE_LAYERS: layer = net[style_layer] bs, height, width, filters = map(lambda i:i.value,layer.get_shape()) size = height * width * filters feats = tf.reshape(layer, (bs, height * width, filters)) feats_T = tf.transpose(feats, perm=[0,2,1]) grams = tf.matmul(feats_T, feats) / size style_gram = style_features[style_layer] style_losses.append(2 * tf.nn.l2_loss(grams - style_gram)/style_gram.size) style_loss = style_weight * functools.reduce(tf.add, style_losses) / batch_size # total variation denoising tv_y_size = _tensor_size(preds[:,1:,:,:]) tv_x_size = _tensor_size(preds[:,:,1:,:]) y_tv = tf.nn.l2_loss(preds[:,1:,:,:] - preds[:,:batch_shape[1]-1,:,:]) x_tv = tf.nn.l2_loss(preds[:,:,1:,:] - preds[:,:,:batch_shape[2]-1,:]) tv_loss = tv_weight*2*(x_tv/tv_x_size + y_tv/tv_y_size)/batch_size loss = content_loss + style_loss + tv_loss # overall loss train_step = tf.train.AdamOptimizer(learning_rate).minimize(loss) sess.run(tf.global_variables_initializer()) import random uid = random.randint(1, 100) print("UID: %s" % uid) for epoch in range(epochs): num_examples = len(content_targets) iterations = 0 while iterations * batch_size < num_examples: start_time = time.time() curr = iterations * batch_size step = curr + batch_size X_batch = np.zeros(batch_shape, dtype=np.float32) for j, img_p in enumerate(content_targets[curr:step]): X_batch[j] = get_img(img_p, (256,256,3)).astype(np.float32) iterations += 1 assert X_batch.shape[0] == batch_size feed_dict = { X_content:X_batch } train_step.run(feed_dict=feed_dict) end_time = time.time() delta_time = end_time - start_time if debug: print("UID: %s, batch time: %s" % (uid, delta_time)) is_print_iter = int(iterations) % print_iterations == 0 if slow: is_print_iter = epoch % print_iterations == 0 is_last = epoch == epochs - 1 and iterations * batch_size >= num_examples should_print = is_print_iter or is_last if should_print: to_get = [style_loss, content_loss, tv_loss, loss, preds] test_feed_dict = { X_content:X_batch } tup = sess.run(to_get, feed_dict = test_feed_dict) _style_loss,_content_loss,_tv_loss,_loss,_preds = tup losses = (_style_loss, _content_loss, _tv_loss, _loss) if slow: _preds = vgg.unprocess(_preds) else: saver = tf.train.Saver() res = saver.save(sess, save_path) yield(_preds, losses, iterations, epoch)