Spaces:
Sleeping
Sleeping
import os | |
import torch as th | |
import torch.multiprocessing as mp | |
import threading as mt | |
import numpy as np | |
import random | |
import ttools | |
import pydiffvg | |
import time | |
def render(canvas_width, canvas_height, shapes, shape_groups, samples=2, | |
seed=None): | |
if seed is None: | |
seed = random.randint(0, 1000000) | |
_render = pydiffvg.RenderFunction.apply | |
scene_args = pydiffvg.RenderFunction.serialize_scene( | |
canvas_width, canvas_height, shapes, shape_groups) | |
img = _render(canvas_width, canvas_height, samples, samples, | |
seed, # seed | |
None, # background image | |
*scene_args) | |
return img | |
def opacityStroke2diffvg(strokes, canvas_size=128, debug=False, relative=True, | |
force_cpu=True): | |
dev = strokes.device | |
if force_cpu: | |
strokes = strokes.to("cpu") | |
# pydiffvg.set_use_gpu(False) | |
# if strokes.is_cuda: | |
# pydiffvg.set_use_gpu(True) | |
"""Rasterize strokes given in (dx, dy, opacity) sequence format.""" | |
bs, nsegs, dims = strokes.shape | |
out = [] | |
start = time.time() | |
for batch_idx, stroke in enumerate(strokes): | |
if relative: # Absolute coordinates | |
all_points = stroke[..., :2].cumsum(0) | |
else: | |
all_points = stroke[..., :2] | |
all_opacities = stroke[..., 2] | |
# Transform from [-1, 1] to canvas coordinates | |
# Make sure points are in canvas | |
all_points = 0.5*(all_points + 1.0) * canvas_size | |
# all_points = th.clamp(0.5*(all_points + 1.0), 0, 1) * canvas_size | |
# Avoid overlapping points | |
eps = 1e-4 | |
all_points = all_points + eps*th.randn_like(all_points) | |
shapes = [] | |
shape_groups = [] | |
for start_idx in range(0, nsegs-1): | |
points = all_points[start_idx:start_idx+2].contiguous().float() | |
opacity = all_opacities[start_idx] | |
num_ctrl_pts = th.zeros(points.shape[0] - 1, dtype=th.int32) | |
width = th.ones(1) | |
path = pydiffvg.Path( | |
num_control_points=num_ctrl_pts, points=points, | |
stroke_width=width, is_closed=False) | |
shapes.append(path) | |
color = th.cat([th.ones(3, device=opacity.device), | |
opacity.unsqueeze(0)], 0) | |
path_group = pydiffvg.ShapeGroup( | |
shape_ids=th.tensor([len(shapes) - 1]), | |
fill_color=None, | |
stroke_color=color) | |
shape_groups.append(path_group) | |
# Rasterize only if there are shapes | |
if shapes: | |
inner_start = time.time() | |
out.append(render(canvas_size, canvas_size, shapes, shape_groups, | |
samples=4)) | |
if debug: | |
inner_elapsed = time.time() - inner_start | |
print("diffvg call took %.2fms" % inner_elapsed) | |
else: | |
out.append(th.zeros(canvas_size, canvas_size, 4, | |
device=strokes.device)) | |
if debug: | |
elapsed = (time.time() - start)*1000 | |
print("rendering took %.2fms" % elapsed) | |
images = th.stack(out, 0).permute(0, 3, 1, 2).contiguous() | |
# Return data on the same device as input | |
return images.to(dev) | |
def stroke2diffvg(strokes, canvas_size=128): | |
"""Rasterize strokes given some sequential data.""" | |
bs, nsegs, dims = strokes.shape | |
out = [] | |
for stroke_idx, stroke in enumerate(strokes): | |
end_of_stroke = stroke[:, 4] == 1 | |
last = end_of_stroke.cpu().numpy().argmax() | |
stroke = stroke[:last+1, :] | |
# stroke = stroke[~end_of_stroke] | |
# TODO: stop at the first end of stroke | |
# import ipdb; ipdb.set_trace() | |
split_idx = stroke[:, 3].nonzero().squeeze(1) | |
# Absolute coordinates | |
all_points = stroke[..., :2].cumsum(0) | |
# Transform to canvas coordinates | |
all_points[..., 0] += 0.5 | |
all_points[..., 0] *= canvas_size | |
all_points[..., 1] += 0.5 | |
all_points[..., 1] *= canvas_size | |
# Make sure points are in canvas | |
all_points[..., :2] = th.clamp(all_points[..., :2], 0, canvas_size) | |
shape_groups = [] | |
shapes = [] | |
start_idx = 0 | |
for count, end_idx in enumerate(split_idx): | |
points = all_points[start_idx:end_idx+1].contiguous().float() | |
if points.shape[0] <= 2: # we need at least 2 points for a line | |
continue | |
num_ctrl_pts = th.zeros(points.shape[0] - 1, dtype=th.int32) | |
width = th.ones(1) | |
path = pydiffvg.Path( | |
num_control_points=num_ctrl_pts, points=points, | |
stroke_width=width, is_closed=False) | |
start_idx = end_idx+1 | |
shapes.append(path) | |
color = th.ones(4, 1) | |
path_group = pydiffvg.ShapeGroup( | |
shape_ids=th.tensor([len(shapes) - 1]), | |
fill_color=None, | |
stroke_color=color) | |
shape_groups.append(path_group) | |
# Rasterize | |
if shapes: | |
# draw only if there are shapes | |
out.append(render(canvas_size, canvas_size, shapes, shape_groups, samples=2)) | |
else: | |
out.append(th.zeros(canvas_size, canvas_size, 4, | |
device=strokes.device)) | |
return th.stack(out, 0).permute(0, 3, 1, 2)[:, :3].contiguous() | |
def line_render(all_points, all_widths, all_alphas, force_cpu=True, | |
canvas_size=32, colors=None): | |
dev = all_points.device | |
if force_cpu: | |
all_points = all_points.to("cpu") | |
all_widths = all_widths.to("cpu") | |
all_alphas = all_alphas.to("cpu") | |
if colors is not None: | |
colors = colors.to("cpu") | |
all_points = 0.5*(all_points + 1.0) * canvas_size | |
eps = 1e-4 | |
all_points = all_points + eps*th.randn_like(all_points) | |
bs, num_segments, _, _ = all_points.shape | |
n_out = 3 if colors is not None else 1 | |
output = th.zeros(bs, n_out, canvas_size, canvas_size, | |
device=all_points.device) | |
scenes = [] | |
for k in range(bs): | |
shapes = [] | |
shape_groups = [] | |
for p in range(num_segments): | |
points = all_points[k, p].contiguous().cpu() | |
num_ctrl_pts = th.zeros(1, dtype=th.int32) | |
width = all_widths[k, p].cpu() | |
alpha = all_alphas[k, p].cpu() | |
if colors is not None: | |
color = colors[k, p] | |
else: | |
color = th.ones(3, device=alpha.device) | |
color = th.cat([color, alpha.view(1,)]) | |
path = pydiffvg.Path( | |
num_control_points=num_ctrl_pts, points=points, | |
stroke_width=width, is_closed=False) | |
shapes.append(path) | |
path_group = pydiffvg.ShapeGroup( | |
shape_ids=th.tensor([len(shapes) - 1]), | |
fill_color=None, | |
stroke_color=color) | |
shape_groups.append(path_group) | |
# Rasterize | |
scenes.append((canvas_size, canvas_size, shapes, shape_groups)) | |
raster = render(canvas_size, canvas_size, shapes, shape_groups, | |
samples=2) | |
raster = raster.permute(2, 0, 1).view(4, canvas_size, canvas_size) | |
alpha = raster[3:4] | |
if colors is not None: # color output | |
image = raster[:3] | |
alpha = alpha.repeat(3, 1, 1) | |
else: | |
image = raster[:1] | |
# alpha compositing | |
image = image*alpha | |
output[k] = image | |
output = output.to(dev) | |
return output, scenes | |
def bezier_render(all_points, all_widths, all_alphas, force_cpu=True, | |
canvas_size=32, colors=None): | |
dev = all_points.device | |
if force_cpu: | |
all_points = all_points.to("cpu") | |
all_widths = all_widths.to("cpu") | |
all_alphas = all_alphas.to("cpu") | |
if colors is not None: | |
colors = colors.to("cpu") | |
all_points = 0.5*(all_points + 1.0) * canvas_size | |
eps = 1e-4 | |
all_points = all_points + eps*th.randn_like(all_points) | |
bs, num_strokes, num_pts, _ = all_points.shape | |
num_segments = (num_pts - 1) // 3 | |
n_out = 3 if colors is not None else 1 | |
output = th.zeros(bs, n_out, canvas_size, canvas_size, | |
device=all_points.device) | |
scenes = [] | |
for k in range(bs): | |
shapes = [] | |
shape_groups = [] | |
for p in range(num_strokes): | |
points = all_points[k, p].contiguous().cpu() | |
# bezier | |
num_ctrl_pts = th.zeros(num_segments, dtype=th.int32) + 2 | |
width = all_widths[k, p].cpu() | |
alpha = all_alphas[k, p].cpu() | |
if colors is not None: | |
color = colors[k, p] | |
else: | |
color = th.ones(3, device=alpha.device) | |
color = th.cat([color, alpha.view(1,)]) | |
path = pydiffvg.Path( | |
num_control_points=num_ctrl_pts, points=points, | |
stroke_width=width, is_closed=False) | |
shapes.append(path) | |
path_group = pydiffvg.ShapeGroup( | |
shape_ids=th.tensor([len(shapes) - 1]), | |
fill_color=None, | |
stroke_color=color) | |
shape_groups.append(path_group) | |
# Rasterize | |
scenes.append((canvas_size, canvas_size, shapes, shape_groups)) | |
raster = render(canvas_size, canvas_size, shapes, shape_groups, | |
samples=2) | |
raster = raster.permute(2, 0, 1).view(4, canvas_size, canvas_size) | |
alpha = raster[3:4] | |
if colors is not None: # color output | |
image = raster[:3] | |
alpha = alpha.repeat(3, 1, 1) | |
else: | |
image = raster[:1] | |
# alpha compositing | |
image = image*alpha | |
output[k] = image | |
output = output.to(dev) | |
return output, scenes | |