Spaces:

vivien
/

depth-aware-caption

Sleeping

Vivien

Links added in the sidebar

03236d3 almost 2 years ago

No virus

6.87 kB

	import numpy as np
	from PIL import ImageDraw, Image, ImageFont
	from transformers import DPTFeatureExtractor, DPTForDepthEstimation
	import torch
	import streamlit as st

	FONTS = [
	"Font: Serif - EBGaramond",
	"Font: Serif - Cinzel",
	"Font: Sans - Roboto",
	"Font: Sans - Lato",
	"Font: Display - Lobster",
	"Font: Display - LilitaOne",
	"Font: Handwriting - GreatVibes",
	"Font: Handwriting - Pacifico",
	"Font: Mono - Inconsolata",
	"Font: Mono - Cutive",
	]


	def hex_to_rgb(hex):
	rgb = []
	for i in (0, 2, 4):
	decimal = int(hex[i : i + 2], 16)
	rgb.append(decimal)
	return tuple(rgb)


	@st.cache(allow_output_mutation=True)
	def load():
	feature_extractor = DPTFeatureExtractor.from_pretrained("Intel/dpt-large")
	model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large")
	return model, feature_extractor


	model, feature_extractor = load()


	def compute_depth(image):
	inputs = feature_extractor(images=image, return_tensors="pt")
	with torch.no_grad():
	outputs = model(**inputs)
	predicted_depth = outputs.predicted_depth
	prediction = torch.nn.functional.interpolate(
	predicted_depth.unsqueeze(1),
	size=image.size[::-1],
	mode="bicubic",
	align_corners=False,
	)
	return prediction.cpu().numpy()[0, 0, :, :]


	def get_mask1(
	shape, x, y, caption, font=None, font_size=0.08, color=(0, 0, 0), alpha=0.8
	):
	img_text = Image.new("RGBA", (shape[1], shape[0]), (0, 0, 0, 0))
	draw = ImageDraw.Draw(img_text)
	font = ImageFont.truetype(font, int(font_size * shape[1]))
	draw.text(
	(x * shape[1], (1 - y) * shape[0]),
	caption,
	fill=(color, int(max(min(1, alpha), 0) 255)),
	font=font,
	)
	text = np.array(img_text)
	mask1 = np.dot(np.expand_dims(text[:, :, -1] / 255, -1), np.ones((1, 3)))
	return text[:, :, :-1], mask1


	def get_mask2(depth_map, depth):
	return np.expand_dims(
	(depth_map[:, :] < depth * np.min(depth_map) + (1 - depth) * np.max(depth_map)),
	-1,
	)


	def add_caption(
	img,
	caption,
	depth_map=None,
	x=0.5,
	y=0.5,
	depth=0.5,
	font_size=50,
	color=(255, 255, 255),
	font="",
	alpha=1,
	):
	text, mask1 = get_mask1(
	img.shape,
	x,
	y,
	caption,
	font=font,
	font_size=font_size,
	color=color,
	alpha=alpha,
	)
	mask2 = get_mask2(depth_map, depth)
	mask = mask1 * np.dot(mask2, np.ones((1, 3)))

	return ((1 - mask) * img + mask * text).astype(np.uint8)


	@st.cache(max_entries=30, show_spinner=False)
	def load_img(uploaded_file):
	if uploaded_file is None:
	img = Image.open("pulp.jpg")
	default = True
	else:
	img = Image.open(uploaded_file)
	if img.size[0] > 800 or img.size[1] > 800:
	if img.size[0] < img.size[1]:
	new_size = (int(800 * img.size[0] / img.size[1]), 800)
	else:
	new_size = (800, int(800 * img.size[1] / img.size[0]))
	img = img.resize(new_size)
	default = False
	return np.array(img), compute_depth(img), default


	def main():
	st.markdown(
	"""
	<style>
	label{
	height: 0px !important;
	min-height: 0px !important;
	margin-bottom: 0px !important;
	}
	</style>
	""",
	unsafe_allow_html=True,
	)

	st.sidebar.markdown(
	"""
	# Depth-aware text addition

	Add text *inside* an image!

	Upload an image, enter some text and adjust the *depth* where you want the text to be displayed. You can also define its location and appearance (font, color, transparency and size).

	Built with [PyTorch](https://pytorch.org/), Intel's [MiDaS model](https://pytorch.org/hub/intelisl_midas_v2/), [Streamlit](https://streamlit.io/), [pillow](https://python-pillow.org/) and inspired by the official [video](https://youtu.be/eTa1jHk1Lxc) of Jenny of Oldstones by Florence + the Machine

	To go further:
	- [blog post](https://vivien000.github.io/blog/journal/adding-text-inside-pictures-and-videos.html)
	- [notebook](https://colab.research.google.com/github/vivien000/depth-aware_captioning/blob/master/Depth_aware_Video_Captioning.ipynb) for videos
	- [examples](https://youtu.be/RtkBplRuWhg?list=PLlPB25tBWqtVhj4Ink8hl9Evc2dlIX4Jh) of videos
	"""
	)

	uploaded_file = st.file_uploader("", type=["jpg", "jpeg"])
	with st.spinner("Analyzing the image - Please wait a few seconds"):
	img, depth_map, default = load_img(uploaded_file)

	if default:
	x0, y0, alpha0, font_size0, depth0, font0 = 0.02, 0.68, 0.99, 0.07, 0.12, 4
	text0 = "Pulp Fiction"
	else:
	x0, y0, alpha0, font_size0, depth0, font0 = 0.1, 0.9, 0.8, 0.08, 0.5, 0
	text0 = "Enter your text here"

	colA, colB, colC = st.columns((13, 1, 1))

	with colA:
	text = st.text_input("", text0)

	with colB:
	st.markdown("Color:")

	with colC:
	color = st.color_picker("", value="#FFFFFF")

	col1, _, col2 = st.columns((4, 1, 4))

	with col1:
	depth = st.select_slider(
	"",
	options=[i / 100 for i in range(101)],
	value=depth0,
	format_func=lambda x: "Foreground"
	if x == 0.0
	else "Background"
	if x == 1.0
	else "",
	)
	x = st.select_slider(
	"",
	options=[i / 100 for i in range(101)],
	value=x0,
	format_func=lambda x: "Left" if x == 0.0 else "Right" if x == 1.0 else "",
	)
	y = st.select_slider(
	"",
	options=[i / 100 for i in range(101)],
	value=y0,
	format_func=lambda x: "Bottom" if x == 0.0 else "Top" if x == 1.0 else "",
	)

	with col2:
	font_size = st.select_slider(
	"",
	options=[0.04 + i / 100 for i in range(0, 17)],
	value=font_size0,
	format_func=lambda x: "Small font"
	if x == 0.04
	else "Large font"
	if x == 0.2
	else "",
	)
	alpha = st.select_slider(
	"",
	options=[i / 100 for i in range(101)],
	value=alpha0,
	format_func=lambda x: "Transparent"
	if x == 0.0
	else "Opaque"
	if x == 1.0
	else "",
	)
	font = st.selectbox("", FONTS, index=font0)

	font = f"fonts/{font[6:]}.ttf"

	captioned = add_caption(
	img,
	text,
	x=x,
	y=y,
	depth=depth,
	depth_map=depth_map,
	font=font,
	font_size=font_size,
	alpha=alpha,
	color=hex_to_rgb(color[1:]),
	)

	st.image(captioned)


	if __name__ == "__main__":
	main()