def _forward(self, z):
bs = z.shape[0]
feats = self.trunk(z)
all_points = self.point_predictor(feats)
all_alphas = self.alpha_predictor(feats)
if self.color_predictor:
all_colors = self.color_predictor(feats)
all_colors = all_colors.view(bs, self.num_strokes, 3)
else:
all_colors = None
all_widths = self.width_predictor(feats)
min_width = self.stroke_width[0]
max_width = self.stroke_width[1]
all_widths = (max_width - min_width) * all_widths + min_width
all_points = all_points.view(bs, self.num_strokes, 2, 2)
output, scenes = rendering.line_render(all_points, all_widths, all_alphas,
colors=all_colors,
canvas_size=self.imsize)
# map to [-1, 1]
output = output*2.0 - 1.0
return output, scenes
def _forward(self, z, hidden_and_cell=None):
steps = self.num_strokes
# z is passed at each step, duplicate it
bs = z.shape[0]
expanded_z = z.unsqueeze(1).repeat(1, steps, 1)
# First step in the RNN
if hidden_and_cell is None:
# Initialize from latent vector
hidden_and_cell = self.hidden_cell_predictor(th.tanh(z))
hidden = hidden_and_cell[:, :self.hidden_size * self.num_layers]
hidden = hidden.view(-1, self.num_layers, self.hidden_size)
hidden = hidden.permute(1, 0, 2).contiguous()
cell = hidden_and_cell[:, self.hidden_size * self.num_layers:]
cell = cell.view(-1, self.num_layers, self.hidden_size)
cell = cell.permute(1, 0, 2).contiguous()
hidden_and_cell = (hidden, cell)
feats, hidden_and_cell = self.lstm(expanded_z, hidden_and_cell)
hidden, cell = hidden_and_cell
feats = feats.reshape(bs * steps, self.hidden_size)
# Construct the chain
end_points = self.point_predictor(feats).view(bs, steps, 1, 2)
start_points = th.cat(
[
# first point is canvas center
th.zeros(bs, 1, 1, 2, device=feats.device),
end_points[:, 1:, :, :]
],
1)
all_points = th.cat([start_points, end_points], 2)
all_alphas = self.alpha_predictor(feats).view(bs, steps)
all_widths = self.width_predictor(feats).view(bs, steps)
min_width = self.stroke_width[0]
max_width = self.stroke_width[1]
all_widths = (max_width - min_width) * all_widths + min_width
output, scenes = rendering.line_render(all_points,
all_widths,
all_alphas,
canvas_size=self.imsize)
# map to [-1, 1]
output = output * 2.0 - 1.0
return output, scenes