drawing-machine/linedraw.py

#!/usr/bin/env python3
"""
Takes a raster image file and vectorizes it to create line art usable by
a pen plotter.
"""
import os
import math
from random import *

from PIL import Image, ImageDraw, ImageOps

no_cv = True
draw_contours = True
draw_hatch = False

try:
	import numpy as np
	import cv2
except:
	print("Cannot import numpy/openCV. Switching to NO_CV mode.")
	no_cv = True


F_Blur = {
	(-2,-2):2,(-1,-2):4,(0,-2):5,(1,-2):4,(2,-2):2,
	(-2,-1):4,(-1,-1):9,(0,-1):12,(1,-1):9,(2,-1):4,
	(-2,0):5,(-1,0):12,(0,0):15,(1,0):12,(2,0):5,
	(-2,1):4,(-1,1):9,(0,1):12,(1,1):9,(2,1):4,
	(-2,2):2,(-1,2):4,(0,2):5,(1,2):4,(2,2):2,
}

F_SobelX = {
	(-1,-1): 1,
	(0,-1):  0,
	(1,-1): -1,
	(-1,0):  2,
	(0,0):   0,
	(1,0):  -2,
	(-1,1):  1,
	(0,1):   0,
	(1,1):  -1}

F_SobelY = {
	(-1,-1): 1,
	(0,-1):  2,
	(1,-1):  1,
	(-1,0):  0,
	(0,0):   0,
	(1,0):   0,
	(-1,1): -1,
	(0,1):  -2,
	(1,1):  -1}

def appmask(image, masks):
	PX = image.load()
	w, h = image.size
	NPX = {}
	for x in range(0, w):
		for y in range(0, h):
			a = [0]*len(masks)
			for i in range(len(masks)):
				for p in masks[i].keys():
					if 0 < x + p[0] < w and 0 < y + p[1] < h:
						a[i] += PX[x + p[0], y + p[1]] * masks[i][p]
				if sum(masks[i].values()) != 0:
					a[i] = a[i] / sum(masks[i].values())
			NPX[x,y] = int(sum([v**2 for v in a])**0.5)
	for x in range(0, w):
		for y in range(0, h):
			PX[x, y] = NPX[x, y]


def distsum(*args):
	"""
	Takes a list of pairs of points, finds the distance between each point
	pair, and returns the sum of all distances.
	"""
	dists = []
	for i in range(1, len(args)):
		a = args[i][0] - args[i-1][0]
		b = args[i][1] - args[i-1][1]
		dist = (a**2 + b**2)**0.5
		dists.append(dist)
	return sum(dists)


def sortlines(lines):
	print("optimizing stroke sequence...")
	clines = lines[:]
	slines = [clines.pop(0)]
	while clines != []:
		x, s, r = None, 1000000, False
		for l in clines:
			d = distsum(l[0], slines[-1][-1])
			dr = distsum(l[-1], slines[-1][-1])
			if d < s:
				x, s, r = l[:], d, False
			if dr < s:
				x, s, r = l[:], s, True

		clines.remove(x)
		if r == True:
			x = x[::-1]
		slines.append(x)
	return slines


def auto_canny(image, sigma=0.33):
	"""
	Automatically determines appropriate upper and lower boundries for the Canny function.
	"""
	med = np.median(image)
	lower = int(max(0, (1.0 - sigma) * med))
	upper = int(min(255, (1.0 + sigma) * med))
	edges = cv2.Canny(image, lower, upper)
	return edges


def find_edges(image):
	print("finding edges...")
	no_cv = True
	if no_cv:
		#appmask(image, [F_Blur])
		appmask(image, [F_SobelX, F_SobelY])
	else:
		image = np.array(image) 
		image = cv2.GaussianBlur(image, (3, 3), 0)
		#image = cv2.Canny(image,100,200)
		image = auto_canny(image)
		image = Image.fromarray(image)
	return image.point(lambda p: p > 128 and 255)  


def getdots(image):
	print("getting contour points...")
	PX = image.load()
	dots = []
	w, h = image.size
	for y in range(h-1):
		row = []
		for x in range(1, w):
			if PX[x, y] == 255:
				if len(row) > 0:
					if x-row[-1][0] == row[-1][-1] + 1:
						row[-1] = (row[-1][0], row[-1][-1] + 1)
					else:
						row.append((x, 0))
				else:
					row.append((x, 0))
		dots.append(row)
	return dots


def connectdots(dots):
	print("connecting contour points...")
	contours = []
	for y in range(len(dots)):
		for x, v in dots[y]:
			if v > -1:
				if y == 0:
					contours.append([(x, y)])
				else:
					closest = -1
					cdist = 100
					for x0, v0 in dots[y-1]:
						if abs(x0 - x) < cdist:
							cdist = abs(x0 - x)
							closest = x0

					if cdist > 3:
						contours.append([(x, y)])
					else:
						found = 0
						for i in range(len(contours)):
							if contours[i][-1] == (closest, y-1):
								contours[i].append((x, y,))
								found = 1
								break
						if found == 0:
							contours.append([(x, y)])
		for c in contours:
			if c[-1][1] < y-1 and len(c) < 4:
				contours.remove(c)
	return contours


def getcontours(image, sc=2):
	print("generating contours...")
	image = find_edges(image)
	image1 = image.copy()
	image2 = image.rotate(-90, expand=True).transpose(Image.FLIP_LEFT_RIGHT)
	dots1 = getdots(image1)
	contours1 = connectdots(dots1)
	dots2 = getdots(image2)
	contours2 = connectdots(dots2)

	for i in range(len(contours2)):
		contours2[i] = [(c[1], c[0]) for c in contours2[i]]	
	contours = contours1 + contours2

	for i in range(len(contours)):
		for j in range(len(contours)):
			if len(contours[i]) > 0 and len(contours[j]) > 0:
				if distsum(contours[j][0], contours[i][-1]) < 8:
					contours[i] = contours[i] + contours[j]
					contours[j] = []

	for i in range(len(contours)):
		contours[i] = [contours[i][j] for j in range(0, len(contours[i]), 8)]


	contours = [c for c in contours if len(c) > 1]

	for i in range(0, len(contours)):
		contours[i] = [(v[0] * sc, v[1] * sc) for v in contours[i]]

	for i in range(0, len(contours)):
		for j in range(0, len(contours[i])):
			contours[i][j] = int(contours[i][j][0] + 10), int(contours[i][j][1] + 10)

	return contours


def hatch(image, sc=16):
	print("hatching...")
	PX = image.load()
	w,h = image.size
	lg1 = []
	lg2 = []
	for x0 in range(w):
		for y0 in range(h):
			x = x0 * sc
			y = y0 * sc
			if PX[x0, y0] > 144:
				pass
				
			elif PX[x0, y0] > 64:
				lg1.append([(x, y+sc/4), (x+sc, y+sc/4)])
			elif PX[x0, y0] > 16:
				lg1.append([(x, y+sc/4), (x+sc, y+sc/4)])
				lg2.append([(x+sc, y), (x, y+sc)])

			else:
				lg1.append([(x, y+sc/4), (x+sc, y+sc/4)])
				lg1.append([(x, y+sc/2 + sc/4), (x+sc, y+sc/2 + sc/4)])
				lg2.append([(x+sc, y), (x, y+sc)])

	lines = [lg1, lg2]
	for k in range(0, len(lines)):
		for i in range(0, len(lines[k])):
			for j in range(0, len(lines[k])):
				if lines[k][i] != [] and lines[k][j] != []:
					if lines[k][i][-1] == lines[k][j][0]:
						lines[k][i] = lines[k][i] + lines[k][j][1:]
						lines[k][j] = []
		lines[k] = [l for l in lines[k] if len(l) > 0]
	lines = lines[0] + lines[1]

	for i in range(0, len(lines)):
		for j in range(0, len(lines[i])):
			lines[i][j] = int(lines[i][j][0] + sc), int(lines[i][j][1] + sc) - j
	return lines


def sketch(path, export_path=None, resolution=1024, hatch_size=16, contour_simplify=2):
	image = Image.open(path)
	if image._getexif().get(274): # "Orientation"
		if image._getexif().get(274) > 4: # modes 5-8 indicate the image was rotated
			image = image.rotate(-90, expand=True)

	image = image.convert("L")
	image = ImageOps.autocontrast(image ,10)

	lines = []
	if draw_contours:
		width = int(resolution/contour_simplify)
		height = int(resolution/contour_simplify*image.height/image.width)
		lines += getcontours(image.resize((width, height)), contour_simplify)

	if draw_hatch:
		width = int(resolution/hatch_size)
		height = int(resolution/hatch_size*image.height/image.width)
		lines += hatch(image.resize((width, height)), hatch_size)

	lines = sortlines(lines)

	export_path = path.replace("received", "converted")
	export_path = os.path.splitext(export_path)[0]
	os.makedirs(os.path.split(export_path)[0], exist_ok=True)
	with open(export_path + ".svg", "w") as file:
		file.write(makePathSvg(lines))
	with open(export_path + ".ngc", "w") as file:
		file.write(makeGcode(lines))

	print(len(lines), "strokes.")
	print("done.")
	return lines


def makePolySvg(lines):
	"""
	Uses the provided set of contours to generate an SVG file and returns it as a string. Contours get
	written as polyline as objects.
	"""
	print("generating svg file...")
	out = '<svg xmlns="http://www.w3.org/2000/svg" version="1.1">\n'
	for l in lines:
		l = ",".join([str(p[0]*0.5)+","+str(p[1]*0.5) for p in l])
		out += '<polyline points="'+l+'" stroke="black" stroke-width="2" fill="none" />\n'
	out += '</svg>'
	return out


def makePathSvg(lines):
	"""
	Uses the provided set of contours to generate an SVG file and returns it as a string. Contours get
	written as path objects.
	"""
	print("generating svg file...")
	out = '<svg xmlns="http://www.w3.org/2000/svg" version="1.1">\n'
	for l in lines:
		l = ",".join([str(p[0]*0.5)+","+str(p[1]*0.5) for p in l])
		out += '<path d="M'+l+'" stroke="black" stroke-width="2" fill="none" />\n'
	out += "</svg>"
	return out


def makeGcode(lines, paper_size="letter"):
	"""
	Converts the provided contour lines into G-code commands and returns them as a string.
	"""
	paper_sizes = {
		"letter": (67.46875, 87.3125),
		"ledger": (87.3125, 139.2903226),
		"max": (90, 140)}
	tot = []
	for line in lines:
		tot += line
	max_x = max([p[0] for p in tot])
	max_y = max([p[1] for p in tot])
	d_x = paper_sizes[paper_size][0] / max_x
	d_y = paper_sizes[paper_size][1] / max_y
	scale = min(d_x, d_y)

	print("generating gcode file...")
	out = "$X\n$32=1\nM03\nF800\nG17 G21 G90 G54\nG01\n\n"
	for line in lines:
		start = line.pop(0)
		out += "S1000\n"
		#out += f"X{start[0]*0.5} Y{start[1]*0.5}\n"
		out += "X" + str(start[0]*scale) + " Y" + str(start[1]*scale) + "\n"
		out += "S0\n"
		for point in line:
			#out += f"X{point[0]*0.5} Y{point[1]*0.5}\n"
			out += "X" + str(point[0]*scale) + " Y" + str(point[1]*scale) + "\n"
		out += "\n"
	out += "S1000\nX0 Y0\nM2\n"
	return out


if __name__ == "__main__":
	import argparse

	parser = argparse.ArgumentParser(
		description="Convert image to vectorized line drawing for plotters.")
	parser.add_argument(
		'-i',
		'--input',
		nargs='?',
		help='Input path')
	parser.add_argument(
		'-o',
		'--output',
		nargs='?',
		help='Output path.')
	parser.add_argument(
		'--no_contour',
		action='store_true',
		help="Don't draw contours.")
	parser.add_argument(
		'--hatch',
		action='store_true',
		help='Enable hatching.')
	parser.add_argument(
		'--no_cv',
		action='store_true',
		help="Don't use OpenCV.")
	parser.add_argument(
		"--resolution",
		default=1024,
		type=int,
		help="Resolution. eg. 512, 1024, 2048")
	parser.add_argument(
		'--contour_simplify',
		default=2,
		type=int,
		help='Level of contour simplification. eg. 1, 2, 3')
	parser.add_argument(
		'--hatch_size',
		default=16,
		type=int,
		help='Patch size of hatches. eg. 8, 16, 32')
	args = parser.parse_args()

	export_path = args.output
	draw_hatch = args.hatch
	draw_contours = not args.no_contour
	no_cv = args.no_cv

	sketch(args.input, args.output, args.resolution, args.contour_simplify, args.hatch_size)