INNER CODE UNIT · Python

output

zuruoke/watermark-removal · main.py:42

            output = tf.reverse(output, [-1])
            output = tf.saturate_cast(output, tf.uint8)
            # load pretrained model
            vars_list = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)
            assign_ops = []
            for var in vars_list:
                vname = var.name
                from_name = vname
                var_value = tf.contrib.framework.load_variable(
                    args.checkpoint_dir, from_name)
                assign_ops.append(tf.assign(var, var_value))
            sess.run(assign_ops)
            print('Model loaded.')
            result = sess.run(output)
            cv2.imwrite(args.output, cv2.cvtColor(
                result[0][:, :, ::-1], cv2.COLOR_BGR2RGB))
            print('image saved to {}'.format(args.output))

View source record →