vincentmolin / ddrawdx

A small differentiable drawing library in JAX

ddrawdx

A small set of differentiable and jittable drawing primitives in JAX. A Canvas consists of a data array with pixel values and a mesh of evaluation points. This mesh can be translated, scaled and rotated. Any of the drawing functions evaluates a smooth function on the mesh to create a "mixing" matrix, and mixes pixels in the image with the color supplied. Currently all smooth transitions are sigmoid based. The sharpness of the transition is not scale invariant, i.e. if one zooms out and draws a shape using the same sharpness this will result in a sharper image

It is a bit janky.

(Outdated) Function signature reference here.

Installation

clone the repo and run pip install . in the root directory.

Examples

ddrawdx/example/example.py

Lines 1 to 19 in e7d075f

	import ddrawdx as drx
	import jax.numpy as jnp
	import jax
	import matplotlib.pyplot as plt
	reddish = jnp.array([0.8, 0.2, 0.1])
	c = drx.canvas(200, 200, "RGB")
	c, r = drx.origin(c)
	def flower(c):
	for s in jnp.arange(0.0, 0.8, 0.05):
	c, _ = drx.rotate(c, 1.2 * s)
	c = drx.fill_rect(c, -0.9 + s, -0.9 + s, 0.9 - s, 0.9 - s, reddish * (1 - s), 400 * (0.88 - s))
	return c
	flower = jax.jit(flower)
	c = flower(c)
	fig, ax = drx.show(c)
	ax.set_title("Box flower")
	fig.savefig("boxflower.png")

ddrawdx/example/example.py

Lines 21 to 34 in e7d075f

	c = drx.canvas(200, 200, "RGB")
	c = drx.fill_rect(c, 0.2, 0.2, 0.8, 0.6, jnp.array([0.0, 1.0, 0.0]))
	ps = jnp.array(
	[[0.1, 0.1], [0.1, 0.5], [0.5, 0.1]]
	)
	c = drx.fill_poly(c, ps, color=jnp.array([1.0, 0.0, 0.0]))
	c, r = drx.translate(c, 0.5, 0.5)
	c = drx.draw_line(c, 0.0, 0.0, 0.25, 0.25)
	c, _ = drx.rotate(c, jnp.pi/2)
	c = drx.draw_line(c, 0.0, 0.0, 0.25, 0.25)
	c = drx.restore(c, r)
	fig, ax = drx.show(c)
	ax.set_title("Clock")
	fig.savefig("clock.png")

ddrawdx/example/example.py

Lines 36 to 40 in e7d075f

	c = drx.canvas(200, 200, "GRAY")
	c = drx.fill_rect(c, 0.3, 0.3, 0.7, 0.4, jnp.array([0.2]))
	fig, ax = drx.show(c)
	ax.set_title("Grayscale")
	fig.savefig("gray.png")