Initial commit

.gitignore

@@ -0,0 +1 @@
+result.jpg

README.md

@@ -1,8 +1,8 @@
 ---
-title: Depth Aware Caption
-emoji: 🐢
-colorFrom: red
-colorTo: indigo
+title: Depth-aware text addition
+emoji: ✍️
+colorFrom: green
+colorTo: blue
 sdk: streamlit
 sdk_version: 1.2.0
 app_file: app.py

app.py

@@ -0,0 +1,266 @@
+import numpy as np
+import PIL
+import torch
+import streamlit as st
+import cv2
+
+DEBUG = False
+if DEBUG:
+    device = torch.device("cpu")
+    model_name = "MiDaS_small"
+else:
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    model_name = "DPT_Large"
+
+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",
+]
+
+CACHE_KWARGS = {
+    "show_spinner": False,
+    "hash_funcs": {torch.nn.parameter.Parameter: lambda _: None},
+    "allow_output_mutation": True,
+    "ttl": 900,
+    "max_entries": 20,
+}
+
+
+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(
+    show_spinner=True,
+    hash_funcs={torch.nn.parameter.Parameter: lambda _: None},
+    allow_output_mutation=True,
+)
+def load(model_type):
+    midas = torch.hub.load("intel-isl/MiDaS", model_type)
+    midas.to(device)
+    _ = midas.eval()
+    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+    if model_type == "DPT_Large" or model_type == "DPT_Hybrid":
+        transform = midas_transforms.dpt_transform
+    else:
+        transform = midas_transforms.small_transform
+    return midas, transform
+
+
+midas, transform = load(model_name)
+
+
+@st.cache(**CACHE_KWARGS)
+def compute_depth(img):
+    with torch.no_grad():
+        prediction = midas(transform(img).to(device))
+
+        prediction = torch.nn.functional.interpolate(
+            prediction.unsqueeze(1),
+            size=img.shape[:2],
+            mode="bicubic",
+            align_corners=False,
+        ).squeeze()
+    return prediction.cpu().numpy()
+
+
+@st.cache(**CACHE_KWARGS)
+def get_mask1(shape, caption, font=None, font_size=0.08, color=(0, 0, 0), alpha=0.8):
+    img_text = PIL.Image.new("RGBA", (shape[1], shape[0]), (0, 0, 0, 0))
+    draw = PIL.ImageDraw.Draw(img_text)
+    font = PIL.ImageFont.truetype(font, int(font_size * img.shape[1]))
+    draw.text(
+        (x * img.shape[1], (1 - y) * img.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)))
+    text = text[:, :, :-1]
+    return text, mask1
+
+
+@st.cache(**CACHE_KWARGS)
+def get_mask2(depth_map, depth):
+    m = np.expand_dims(
+        (depth_map[:, :] < depth * np.min(depth_map) + (1 - depth) * np.max(depth_map)),
+        -1,
+    )
+    return np.dot(m, np.ones((1, 3)))
+
+
+@st.cache(**CACHE_KWARGS)
+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,
+):
+    if depth_map is None:
+        depth_map = compute_depth(img)
+    text, mask1 = get_mask1(
+        img.shape, caption, font=font, font_size=font_size, color=color, alpha=alpha
+    )
+    mask2 = get_mask2(depth_map, depth)
+    mask = mask1 * mask2
+
+    return ((1 - mask) * img + mask * text).astype(np.uint8)
+
+
+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 
+"""
+)
+
+uploaded_file = st.file_uploader("", type=["jpg", "jpeg"])
+
+
+@st.cache(**CACHE_KWARGS)
+def load_img(uploaded_file):
+    if uploaded_file is None:
+        img = np.array(PIL.Image.open("pulp.jpg"))
+        default = True
+    else:
+        img = np.array(PIL.Image.open(uploaded_file))
+        if img.shape[0] > 800 or img.shape[1] > 800:
+            if img.shape[0] < img.shape[1]:
+                new_size = (800, int(800 * img.shape[0] / img.shape[1]))
+            else:
+                new_size = (int(800 * img.shape[1] / img.shape[0]), 800)
+            img = cv2.resize(img, dsize=new_size, interpolation=cv2.INTER_CUBIC)
+        default = False
+    depth_map = compute_depth(img)
+    return img, depth_map, default
+
+
+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.23, 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,
+    depth_map=depth_map,
+    x=x,
+    y=y,
+    depth=depth,
+    font=font,
+    font_size=font_size,
+    alpha=alpha,
+    color=hex_to_rgb(color[1:]),
+)
+
+st.image(captioned)
+
+PIL.Image.fromarray(captioned).save("result.jpg")
+with open("result.jpg", "rb") as file:
+    btn = st.download_button(
+        label="Download image", data=file, file_name="result.jpg", mime="image/jpeg"
+    )

packages.txt

@@ -0,0 +1 @@
+libgl1

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+torch
+timm
+pillow
+opencv-python