https://github.com/MarkMoHR/virtual_sketching
Raw File
Tip revision: 958efe45a9120b9d467ba7701efba28c11e38f8f authored by Your Name on 21 August 2021, 08:13:15 UTC
Added bilibili links
Tip revision: 958efe4
NeuralRenderer.py
import tensorflow as tf


class RasterUnit(object):
    def __init__(self,
                 raster_size,
                 input_params,  # (N, 10)
                 reuse=False):
        self.raster_size = raster_size
        self.input_params = input_params

        with tf.variable_scope("raster_unit", reuse=reuse):
            self.build_unit()

    def build_unit(self):
        x = self.input_params  # (N, 10)
        x = self.fully_connected(x, 10, 512, scope='fc1')  # (N, 512)
        x = tf.nn.relu(x)
        x = self.fully_connected(x, 512, 1024, scope='fc2')  # (N, 1024)
        x = tf.nn.relu(x)
        x = self.fully_connected(x, 1024, 2048, scope='fc3')  # (N, 2048)
        x = tf.nn.relu(x)
        x = self.fully_connected(x, 2048, 4096, scope='fc4')  # (N, 4096)
        x = tf.nn.relu(x)
        x = tf.reshape(x, (-1, 16, 16, 16))  # (N, 16, 16, 16)
        x = tf.transpose(x, (0, 2, 3, 1))

        x = self.conv2d(x, 32, 3, 1, scope='conv1')  # (N, 16, 16, 32)
        x = tf.nn.relu(x)
        x = self.conv2d(x, 32, 3, 1, scope='conv2')  # (N, 16, 16, 32)
        x = self.pixel_shuffle(x, upscale_factor=2)  # (N, 32, 32, 8)

        x = self.conv2d(x, 16, 3, 1, scope='conv3')  # (N, 32, 32, 16)
        x = tf.nn.relu(x)
        x = self.conv2d(x, 16, 3, 1, scope='conv4')  # (N, 32, 32, 16)
        x = self.pixel_shuffle(x, upscale_factor=2)  # (N, 64, 64, 4)

        x = self.conv2d(x, 8, 3, 1, scope='conv5')  # (N, 64, 64, 8)
        x = tf.nn.relu(x)
        x = self.conv2d(x, 4, 3, 1, scope='conv6')  # (N, 64, 64, 4)
        x = self.pixel_shuffle(x, upscale_factor=2)  # (N, 128, 128, 1)
        x = tf.sigmoid(x)

        # (N, 128, 128), [0.0-stroke, 1.0-BG]
        self.stroke_image = 1.0 - tf.reshape(x, (-1, self.raster_size, self.raster_size))

    def conv2d(self, input_tensor, out_channels, kernel_size, stride, scope, reuse=False):
        with tf.variable_scope(scope, reuse=reuse):
            output_tensor = tf.layers.conv2d(input_tensor, out_channels, kernel_size=kernel_size,
                                             strides=(stride, stride),
                                             padding="same", kernel_initializer=tf.keras.initializers.he_normal())
            return output_tensor

    def fully_connected(self, input_tensor, in_dim, out_dim, scope, reuse=False):
        with tf.variable_scope(scope, reuse=reuse):
            weight = tf.get_variable("weight", [in_dim, out_dim], dtype=tf.float32,
                                     initializer=tf.random_normal_initializer())
            bias = tf.get_variable("bias", [out_dim], dtype=tf.float32,
                                   initializer=tf.random_normal_initializer())
            output_tensor = tf.matmul(input_tensor, weight) + bias
            return output_tensor

    def pixel_shuffle(self, input_tensor, upscale_factor):
        params_shape = input_tensor.get_shape()
        n, h, w, c = params_shape
        input_tensor_proc = tf.reshape(input_tensor, (n, h, w, c // 4, 4))
        input_tensor_proc = tf.transpose(input_tensor_proc, (0, 1, 2, 4, 3))
        input_tensor_proc = tf.reshape(input_tensor_proc, (n, h, w, -1))
        output_tensor = tf.depth_to_space(input_tensor_proc, block_size=upscale_factor)
        return output_tensor


class NeuralRasterizor(object):
    def __init__(self,
                 raster_size,
                 seq_len,
                 position_format='abs',
                 raster_padding=10,
                 strokes_format=3):
        self.raster_size = raster_size
        self.seq_len = seq_len
        self.position_format = position_format
        self.raster_padding = raster_padding
        self.strokes_format = strokes_format

        assert position_format in ['abs', 'rel']

    def raster_func_abs(self, input_data, raster_seq_len=None):
        """
        x and y in absolute position.
        :param input_data: (N, seq_len, 10): [x0, y0, x1, y1, x2, y2, r0, r2, w0, w2]. All in [0.0, 1.0]
        :return:
        """
        seq_len = raster_seq_len if raster_seq_len is not None else self.seq_len

        raster_params = tf.transpose(input_data, [1, 0, 2])  # (seq_len, N, 10)

        seq_stroke_images = tf.map_fn(self.stroke_drawer_with_raster_unit, raster_params,
                                      parallel_iterations=32)  # (seq_len, N, raster_size, raster_size)
        seq_stroke_images = tf.transpose(seq_stroke_images, [1, 2, 3, 0])
        # (N, raster_size, raster_size, seq_len), [0.0-stroke, 1.0-BG]

        filter_seq_stroke_images = 1.0 - seq_stroke_images
        # (N, raster_size, raster_size, seq_len), [0.0-BG, 1.0-stroke]

        # stacking
        stroke_images_unclip = tf.reduce_sum(filter_seq_stroke_images, axis=-1)  # (N, raster_size, raster_size)
        stroke_images = tf.clip_by_value(stroke_images_unclip, 0.0, 1.0)  # [0.0-BG, 1.0-stroke]
        return stroke_images, stroke_images_unclip, seq_stroke_images

    def stroke_drawer_with_raster_unit(self, params_batch):
        """
        Convert two points into a raster stroke image with RasterUnit.
        :param params_batch: (N, 10)
        :return: (N, raster_size, raster_size)
        """
        raster_unit = RasterUnit(
            raster_size=self.raster_size,
            input_params=params_batch,
            reuse=tf.AUTO_REUSE
        )
        stroke_image = raster_unit.stroke_image  # (N, raster_size, raster_size), [0.0-stroke, 1.0-BG]
        return stroke_image


class NeuralRasterizorStep(object):
    def __init__(self,
                 raster_size,
                 position_format='abs'):
        self.raster_size = raster_size
        self.position_format = position_format

        assert position_format in ['abs', 'rel']

    def raster_func_stroke_abs(self, input_data):
        """
        x and y in absolute position.
        :param input_data: (N, 8): [x0, y0, x1, y1, x2, y2, r0, r2]. All in [0.0, 1.0]
        :return:
        """
        w_in = tf.ones(shape=(input_data.shape[0], 2), dtype=tf.float32)
        raster_params = tf.concat([input_data, w_in], axis=-1)  # (N, 10)
        stroke_image = self.stroke_drawer_with_raster_unit(raster_params)  # (N, raster_size, raster_size), [0.0-stroke, 1.0-BG]
        stroke_image = 1.0 - stroke_image  # [0.0-BG, 1.0-stroke]

        return stroke_image

    def mask_ending_state(self, input_states):
        """
        Mask the ending state to be 1
        :param input_states: (N, seq_len, 1) in offset manner
        :param seq_len:
        :return:
        """
        ending_state_accu = tf.cumsum(input_states, axis=1)  # (N, seq_len, 1)
        ending_state_clip = tf.clip_by_value(ending_state_accu, 0.0, 1.0)  # (N, seq_len, 1)
        return ending_state_clip

    def stroke_drawer_with_raster_unit(self, params_batch):
        """
        Convert two points into a raster stroke image with RasterUnit.
        :param params_batch: (N, 10)
        :return: (N, raster_size, raster_size)
        """
        raster_unit = RasterUnit(
            raster_size=self.raster_size,
            input_params=params_batch,
            reuse=tf.AUTO_REUSE
        )
        stroke_image = raster_unit.stroke_image  # (N, raster_size, raster_size), [0.0-stroke, 1.0-BG]
        return stroke_image
back to top