animation/fourcell/apply_image.py

import argparse
import sys
import cv2
import numpy as np
from json import load
from os.path import exists, basename
from common import image_resize, display, read_json

holeConstant = .0156862745 # 160/10200
fontConstant = 96
markConstant = 5
linesConstant = 150

def generate_text (text, fontSize, rotation) :
	rgb_color=(255,255,255)
	color = tuple(reversed(rgb_color))

	(w, h), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_PLAIN, fontSize, 2)
	image = np.zeros((h + baseline, w, 3), np.uint8)
	image[:] = color
	image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
	image = cv2.putText(image, text, (0, h + round(baseline/2)), cv2.FONT_HERSHEY_PLAIN, fontSize, [0, 0, 0], 2, cv2.LINE_AA)

	if rotation == 90 :
		image = cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)
	elif rotation == -90 :
		image = cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)

	return image

def place_text (output, text_image, i, points) :

	height, width = text_image.shape[:2]
	halfW = int(round(width / 2))
	halfH = int(round(height / 2))

	if i == 0 or i == 1 :
		y = int(round((points['0']['y'] + points['1']['y']) / 2))
		x = int(round((points['0']['x'] + points['1']['x']) / 2)) - width
	elif i == 2 or i == 3 :
		y = int(round((points['0']['y'] + points['1']['y']) / 2))
		x = int(round((points['0']['x'] + points['1']['x']) / 2)) + width

	output[y-halfH:y-halfH+height, x-halfW:x-halfW+width] = text_image

	return output

def to_rotation (i) :
	if i == 0 or i == 1 :
		return -90
	if i == 2 or i == 3 :
		return 90

def to_text (fileName, project_name) :
	parts = fileName.split('.')
	name = parts[0]
	if 'guide_' in name :
		return project_name
	parts = name.split('_')
	return parts[len(parts)-1]

# use top left, top right, bottom left
def apply_image_to_points (original, target, points) :
	rows, cols, ch = original.shape
	ir, ic, ich = target.shape

	print('Using points: ')
	print(points)

	atPts = np.float32(points)
	targetPts = np.float32([[0, 0], [ic, 0], [0, ir]])

	M = cv2.getAffineTransform(targetPts, atPts)
	dst = cv2.warpAffine(target, M, (cols, rows),
		flags = cv2.INTER_CUBIC, 
		borderMode = cv2.BORDER_CONSTANT, 
		borderValue = [0, 0, 0, 0])

	mask = get_mask(dst)
	return combine_images(original, dst, mask)

def create_blank(width, height):
	rgb_color=(255,255,255)
	image = np.zeros((height, width, 3), np.uint8)
	color = tuple(reversed(rgb_color))
	image[:] = color
	#return image
	return cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)

# get points 0, 1 and 3 = top left, top right and bottom left
def to_points (d) :
	return [[d['0']['x'], d['0']['y']], [d['1']['x'], d['1']['y']], [d['3']['x'], d['3']['y']]]

def get_mask (img) :
	lower = np.array([0,0,0,0])
	upper= np.array([0,0,0,0])
	mask = cv2.inRange(img, lower, upper)
	return cv2.bitwise_not(mask)

def combine_images (bg, fg, mask) :
	fk = cv2.bitwise_or(fg, fg, mask=mask)
	mask = cv2.bitwise_not(mask)
	bk = cv2.bitwise_or(bg, bg, mask=mask)
	return cv2.bitwise_or(fk, bk)

def draw_registration_marks (output, points, registrationMark, lines) :
	keys = list(points.keys())

	for key in keys:
		output = draw_registration_mark(output, points[key], registrationMark, lines)

	return output

def draw_registration_mark (output, point, registrationMark, lines) :
	half = int(round(registrationMark / 2))
	topY = point['y'] - half
	bottomY = point['y'] + half
	leftX = point['x'] - half
	rightX = point['x'] + half

	cv2.line(output, (leftX, point['y'],), (rightX, point['y'],), [0, 0, 0], lines) 
	cv2.line(output, (point['x'], topY,), (point['x'], bottomY,), [0, 0, 0], lines)

	return output

def apply_image (args) : 
	if not exists(args.template) :
		print('Calibration template does not exist, please provide one that does')
		exit(3)

	if len(args.images) < 1:
		print('Please provide at least one image to apply')
		exit(4)

	if len(args.images) > 4:
		print('Please provide maximum four images to apply')
		exit(5)

	for img in args.images:
		if not exists(img) :
			print(f'Image {img} does not exist, exiting...')
			exit(6)

	print('Using images: ')
	for img in args.images:
		print(f' -> {basename(img)}')

	tmpl = read_json(args.template)

	print(f"Image size {tmpl['width']}x{tmpl['height']}")

	hole = round(tmpl['width'] * holeConstant)
	blank = create_blank(tmpl['width'], tmpl['height'])
	dpi = round(tmpl['width'] / 8.5)
	fontSize = round(dpi / fontConstant)
	registrationMark = int(round(dpi / markConstant))
	lines = int(round(dpi / linesConstant))
	#blank = cv2.imread(args.template.replace('.calibration.json', ''))
	output = blank.copy()


	hp = tmpl['holePunches']
	for i in hp:
		print(hp[i])
		cv2.circle(output, (hp[i]['x'], hp[i]['y'],), hole, (200, 200, 200,), -1)

	for i in range(0, len(args.images)) :
		print(args.images[i])
		frame = cv2.imread(args.images[i])
		frame = cv2.cvtColor(frame, cv2.COLOR_RGB2RGBA)
		#cv2.imwrite(f'test{i}.png', output)
		output = apply_image_to_points(output, frame, to_points(tmpl[f'{i}']))
		rotation = to_rotation(i)
		text = to_text(basename(args.images[i]), args.name)
		text_image = generate_text(text, fontSize, rotation)
		output = place_text(output, text_image, i, tmpl[f'{i}'])
		output = draw_registration_marks(output, tmpl[f'{i}'], registrationMark, lines)
		print(f'Applied {basename(args.images[i])}')

	output = cv2.cvtColor(output, cv2.COLOR_BGRA2BGR)
	#display(output)
	cv2.imwrite(args.destination, output)
	##print(f'Wrote {args.destination}')

def main () :
	parser = argparse.ArgumentParser()
	parser.add_argument('name',  help='Name of project')
	parser.add_argument('destination',  help='File to output image to')
	parser.add_argument('template', help='Template file to use')
	parser.add_argument('images', nargs=argparse.REMAINDER)
	args = parser.parse_args()

	apply_image(args)

main()