import gradio as gr
from rembg import remove
from PIL import ImageOps, ImageDraw

def process_image(foreImg, flip, backImg, left, top, method, hscale, vscale, width, height, isLoc):
	if foreImg is None:
		return None

	imwidth, imheight = foreImg.size
	if not backImg is None:
		if method == 'horizontal scale factor':
			bgwidth, _ = backImg.size
			scale = bgwidth * hscale / imwidth
			imwidth = int(imwidth * scale)
			imheight = int(imheight * scale)
		elif method == 'vertical scale factor':
			_, bgheight = backImg.size
			scale = bgheight * vscale / imheight
			imwidth = int(imwidth * scale)
			imheight = int(imheight * scale)
		else:
			if width > 0:
				imwidth = width
			if height > 0:
				imheight = height

	if isLoc:
		if backImg is None:
			return None
		draw = ImageDraw.Draw(backImg)
		draw.rectangle([(left, top), (imwidth + left, imheight + top)], outline=(255, 0, 0), width=4)
		return backImg

	image = remove(foreImg)
	if flip:
		image = ImageOps.mirror(image)
	image = image.resize((imwidth, imheight))

	if backImg is None:
		return image

	backImg.paste(image, (left, top), image)
	return backImg

app = gr.Interface(
	title='Combine Images',
	fn=process_image,
	inputs=[
		gr.Image(label='foreground', type='pil'),
		gr.Checkbox(label='flip left and right'),
		gr.Image(label='background', type='pil'),
		gr.Slider(maximum=4000, step=1, label='left'),
		gr.Slider(maximum=4000, step=1, label='top'),
		gr.Radio(['horizontal scale factor', 'vertical scale factor', 'size (width and height)'], label='size specification method', value='horizontal scale factor'),
		gr.Slider(minimum=0.01, maximum=1.0, step=0.01, label='horizontal scale factor', value=0.5),
		gr.Slider(minimum=0.01, maximum=1.0, step=0.01, label='vertical scale factor', value=0.5),
		gr.Slider(maximum=4000, step=1, label='width'),
		gr.Slider(maximum=4000, step=1, label='height'),
		gr.Checkbox(label='check location only'),
	],
	outputs='image',
	allow_flagging='never',
	concurrency_limit=20,
	examples=[['examples/foreground.jpg', False, 'examples/background.jpg', 720, 540, 'size (width and height)', 0.5, 0.5, 256, 256, False]],
	#cache_examples=False
)
app.launch()