Switch to transformers

.gitignore

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

app.py

@@ -1,18 +1,9 @@
 import numpy as np
-import PIL
+from PIL import ImageDraw, Image, ImageFont
+from transformers import DPTFeatureExtractor, DPTForDepthEstimation
 import torch
 import streamlit as st
 
-device = torch.device("cpu")
-
-DEBUG = False
-if DEBUG:
-    cache_kwargs = {"max_entries": 30}
-    model_name = "MiDaS_small"
-else:
-    cache_kwargs = {"show_spinner": False, "max_entries": 30}
-    model_name = "DPT_Large"
-
 FONTS = [
     "Font: Serif - EBGaramond",
     "Font: Serif - Cinzel",
@@ -35,41 +26,36 @@ def hex_to_rgb(hex):
     return tuple(rgb)
 
 
-@st.experimental_singleton
-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
+@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
 
 
-midas, transform = load(model_name)
+model, feature_extractor = load()
 
 
-@st.experimental_memo(**cache_kwargs)
-def compute_depth(img):
+def compute_depth(image):
+    inputs = feature_extractor(images=image, return_tensors="pt")
     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()
+        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 = 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 * shape[1]))
+    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,
@@ -78,8 +64,7 @@ def get_mask1(
     )
     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
+    return text[:, :, :-1], mask1
 
 
 def get_mask2(depth_map, depth):
@@ -101,8 +86,6 @@ def add_caption(
     font="",
     alpha=1,
 ):
-    if depth_map is None:
-        depth_map = compute_depth(img)
     text, mask1 = get_mask1(
         img.shape,
         x,
@@ -119,13 +102,13 @@ def add_caption(
     return ((1 - mask) * img + mask * text).astype(np.uint8)
 
 
-@st.experimental_memo(**cache_kwargs)
+@st.cache(max_entries=30, show_spinner=False)
 def load_img(uploaded_file):
     if uploaded_file is None:
-        img = PIL.Image.open("pulp.jpg")
+        img = Image.open("pulp.jpg")
         default = True
     else:
-        img = PIL.Image.open(uploaded_file)
+        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)
@@ -133,7 +116,7 @@ def load_img(uploaded_file):
                 new_size = (800, int(800 * img.size[1] / img.size[0]))
             img = img.resize(new_size)
         default = False
-    return np.array(img), default
+    return np.array(img), compute_depth(img), default
 
 
 def main():
@@ -163,13 +146,11 @@ def main():
     )
 
     uploaded_file = st.file_uploader("", type=["jpg", "jpeg"])
-
-    img, default = load_img(uploaded_file)
-
-    del uploaded_file
+    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.23, 4
+        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
@@ -243,6 +224,7 @@ def main():
         x=x,
         y=y,
         depth=depth,
+        depth_map=depth_map,
         font=font,
         font_size=font_size,
         alpha=alpha,
@@ -251,14 +233,6 @@ def main():
 
     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"
-    #    )
-
-    del captioned, img
-
 
 if __name__ == "__main__":
     main()

packages.txt

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

requirements.txt

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