Initial commit

.gitignore

@@ -0,0 +1,124 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

.streamlit/config.toml

@@ -0,0 +1,6 @@
+[server]
+maxUploadSize = 10
+
+[theme]
+base="light"
+primaryColor="#0074ff"

Dockerfile

@@ -0,0 +1,9 @@
+FROM pytorch/pytorch:latest
+
+WORKDIR /app
+
+COPY . .
+
+RUN pip install -r requirements.txt
+
+CMD [ "streamlit", "run", "app.py" ]

README.md

@@ -5,6 +5,7 @@ colorFrom: pink
 colorTo: purple
 sdk: streamlit
 sdk_version: 1.2.0
+python_version: 3.9.5
 app_file: app.py
 pinned: false
 ---

app.py

@@ -0,0 +1,126 @@
+# App code based on: https://github.com/petergro-hub/ComicInpainting
+# Model based on: https://github.com/saic-mdal/lama
+
+import numpy as np
+import pandas as pd
+import streamlit as st
+import os
+from datetime import datetime
+from PIL import Image
+from streamlit_drawable_canvas import st_canvas
+from io import BytesIO
+from copy import deepcopy
+
+from src.core import process_inpaint
+
+
+def image_download_button(pil_image, filename: str, fmt: str, label="Download"):
+    if fmt not in ["jpg", "png"]:
+        raise Exception(f"Unknown image format (Available: {fmt} - case sensitive)")
+    
+    pil_format = "JPEG" if fmt == "jpg" else "PNG"
+    file_format = "jpg" if fmt == "jpg" else "png"
+    mime = "image/jpeg" if fmt == "jpg" else "image/png"
+    
+    buf = BytesIO()
+    pil_image.save(buf, format=pil_format)
+    
+    return st.download_button(
+        label=label,
+        data=buf.getvalue(),
+        file_name=f'{filename}.{file_format}',
+        mime=mime,
+    )
+
+
+
+if "button_id" not in st.session_state:
+    st.session_state["button_id"] = ""
+if "color_to_label" not in st.session_state:
+    st.session_state["color_to_label"] = {}
+
+if 'reuse_image' not in st.session_state:
+    st.session_state.reuse_image = None
+def set_image(img):
+    st.session_state.reuse_image = img
+
+st.title("AI Photo Object Removal")
+
+st.image(open("assets/demo.png", "rb").read())
+
+st.markdown(
+    """
+    So you want to remove an object in your photo? You don't need to learn photo editing skills.
+    **Just draw over the parts of the image you want to remove, then our AI will remove them.**
+    """
+)
+uploaded_file = st.file_uploader("Choose image", accept_multiple_files=False, type=["png", "jpg", "jpeg"])
+
+if uploaded_file is not None:
+    
+    if st.session_state.reuse_image is not None:
+        img_input = Image.fromarray(st.session_state.reuse_image)
+    else:
+        bytes_data = uploaded_file.getvalue()
+        img_input = Image.open(BytesIO(bytes_data)).convert("RGBA")
+
+    stroke_width = st.slider("Brush size", 1, 100, 50)
+
+    st.write("Now draw (brush) the part of image that you want to remove.")
+    
+    # Canvas size logic
+    canvas_bg = deepcopy(img_input)
+    aspect_ratio = canvas_bg.width / canvas_bg.height
+    streamlit_width = 720
+    
+    # Max width is 720. Resize the height to maintain its aspectratio.
+    if canvas_bg.width > streamlit_width:
+        canvas_bg = canvas_bg.resize((streamlit_width, int(streamlit_width / aspect_ratio)))
+    
+    canvas_result = st_canvas(
+        stroke_color="rgba(255, 0, 255, 1)",
+        stroke_width=stroke_width,
+        background_image=canvas_bg,
+        width=canvas_bg.width,
+        height=canvas_bg.height,
+        drawing_mode="freedraw",
+        key="compute_arc_length", 
+    )
+    
+    if canvas_result.image_data is not None:
+        im = np.array(Image.fromarray(canvas_result.image_data).resize(img_input.size))
+        background = np.where(
+            (im[:, :, 0] == 0) & 
+            (im[:, :, 1] == 0) & 
+            (im[:, :, 2] == 0)
+        )
+        drawing = np.where(
+            (im[:, :, 0] == 255) & 
+            (im[:, :, 1] == 0) & 
+            (im[:, :, 2] == 255)
+        )
+        im[background]=[0,0,0,255]
+        im[drawing]=[0,0,0,0] #RGBA
+        
+        reuse = False
+        
+        if st.button('Submit'):
+            
+            with st.spinner("AI is doing the magic!"):
+                output = process_inpaint(np.array(img_input), np.array(im)) #TODO Put button here
+                img_output = Image.fromarray(output).convert("RGB")
+            
+            st.write("AI has finished the job!")
+            st.image(img_output)
+            # reuse = st.button('Edit again (Re-use this image)', on_click=set_image, args=(inpainted_img, ))
+            
+            uploaded_name = os.path.splitext(uploaded_file.name)[0]
+            image_download_button(
+                pil_image=img_output,
+                filename=uploaded_name,
+                fmt="jpg",
+                label="Download Image"
+            )
+                
+            st.info("**TIP**: If the result is not perfect, you can download it then "
+                    "upload then remove the artifacts.")

docker-compose.yml

@@ -0,0 +1,14 @@
+---
+version: '3'
+services:
+  st-remove-photo-object:
+    build: .
+    container_name: st-remove-photo-object
+    restart: unless-stopped
+    ports:
+    - 51003:8501
+    volumes:
+    - .:/app
+    environment:
+    - TZ=Asia/Jakarta
+    # command: streamlit run sdc.py

requirements.txt

@@ -0,0 +1,25 @@
+torch
+torchvision
+numpy
+opencv-python-headless
+matplotlib
+streamlit
+streamlit-drawable-canvas
+pyyaml
+tqdm
+easydict==1.9.0
+scikit-image==0.17.2
+scikit-learn==0.24.2
+opencv-python
+tensorflow
+joblib
+pandas
+albumentations==0.5.2
+hydra-core==1.1.0
+pytorch-lightning==1.2.9
+tabulate
+kornia==0.5.0
+webdataset
+packaging
+scikit-learn==0.24.2
+wldhx.yadisk-direct

src/core.py

@@ -0,0 +1,466 @@
+import base64
+import json
+import os
+import re
+import time
+import uuid
+from io import BytesIO
+from pathlib import Path
+import cv2   
+
+# For inpainting
+
+import numpy as np
+import pandas as pd
+import streamlit as st
+from PIL import Image
+from streamlit_drawable_canvas import st_canvas
+
+
+import argparse
+import io
+import multiprocessing
+from typing import Union
+
+import torch
+
+try:
+    torch._C._jit_override_can_fuse_on_cpu(False)
+    torch._C._jit_override_can_fuse_on_gpu(False)
+    torch._C._jit_set_texpr_fuser_enabled(False)
+    torch._C._jit_set_nvfuser_enabled(False)
+except:
+    pass
+
+from src.helper import (
+    download_model,
+    load_img,
+    norm_img,
+    numpy_to_bytes,
+    pad_img_to_modulo,
+    resize_max_size,
+)
+
+NUM_THREADS = str(multiprocessing.cpu_count())
+
+os.environ["OMP_NUM_THREADS"] = NUM_THREADS
+os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS
+os.environ["MKL_NUM_THREADS"] = NUM_THREADS
+os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS
+os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS
+if os.environ.get("CACHE_DIR"):
+    os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"]
+
+#BUILD_DIR = os.environ.get("LAMA_CLEANER_BUILD_DIR", "./lama_cleaner/app/build")
+
+# For Seam-carving
+
+from scipy import ndimage as ndi
+
+SEAM_COLOR = np.array([255, 200, 200])    # seam visualization color (BGR)
+SHOULD_DOWNSIZE = True                    # if True, downsize image for faster carving
+DOWNSIZE_WIDTH = 500                      # resized image width if SHOULD_DOWNSIZE is True
+ENERGY_MASK_CONST = 100000.0              # large energy value for protective masking
+MASK_THRESHOLD = 10                       # minimum pixel intensity for binary mask
+USE_FORWARD_ENERGY = True                 # if True, use forward energy algorithm
+
+device = torch.device("cpu")
+model_path = "./assets/big-lama.pt"
+model = torch.jit.load(model_path, map_location="cpu")
+model = model.to(device)
+model.eval()
+
+
+########################################
+# UTILITY CODE
+########################################
+
+
+def visualize(im, boolmask=None, rotate=False):
+    vis = im.astype(np.uint8)
+    if boolmask is not None:
+        vis[np.where(boolmask == False)] = SEAM_COLOR
+    if rotate:
+        vis = rotate_image(vis, False)
+    cv2.imshow("visualization", vis)
+    cv2.waitKey(1)
+    return vis
+
+def resize(image, width):
+    dim = None
+    h, w = image.shape[:2]
+    dim = (width, int(h * width / float(w)))
+    image = image.astype('float32')
+    return cv2.resize(image, dim)
+
+def rotate_image(image, clockwise):
+    k = 1 if clockwise else 3
+    return np.rot90(image, k)    
+
+
+########################################
+# ENERGY FUNCTIONS
+########################################
+
+def backward_energy(im):
+    """
+    Simple gradient magnitude energy map.
+    """
+    xgrad = ndi.convolve1d(im, np.array([1, 0, -1]), axis=1, mode='wrap')
+    ygrad = ndi.convolve1d(im, np.array([1, 0, -1]), axis=0, mode='wrap')
+    
+    grad_mag = np.sqrt(np.sum(xgrad**2, axis=2) + np.sum(ygrad**2, axis=2))
+
+    # vis = visualize(grad_mag)
+    # cv2.imwrite("backward_energy_demo.jpg", vis)
+
+    return grad_mag
+
+def forward_energy(im):
+    """
+    Forward energy algorithm as described in "Improved Seam Carving for Video Retargeting"
+    by Rubinstein, Shamir, Avidan.
+    Vectorized code adapted from
+    https://github.com/axu2/improved-seam-carving.
+    """
+    h, w = im.shape[:2]
+    im = cv2.cvtColor(im.astype(np.uint8), cv2.COLOR_BGR2GRAY).astype(np.float64)
+
+    energy = np.zeros((h, w))
+    m = np.zeros((h, w))
+    
+    U = np.roll(im, 1, axis=0)
+    L = np.roll(im, 1, axis=1)
+    R = np.roll(im, -1, axis=1)
+    
+    cU = np.abs(R - L)
+    cL = np.abs(U - L) + cU
+    cR = np.abs(U - R) + cU
+    
+    for i in range(1, h):
+        mU = m[i-1]
+        mL = np.roll(mU, 1)
+        mR = np.roll(mU, -1)
+        
+        mULR = np.array([mU, mL, mR])
+        cULR = np.array([cU[i], cL[i], cR[i]])
+        mULR += cULR
+
+        argmins = np.argmin(mULR, axis=0)
+        m[i] = np.choose(argmins, mULR)
+        energy[i] = np.choose(argmins, cULR)
+    
+    # vis = visualize(energy)
+    # cv2.imwrite("forward_energy_demo.jpg", vis)     
+        
+    return energy
+
+########################################
+# SEAM HELPER FUNCTIONS
+######################################## 
+
+def add_seam(im, seam_idx):
+    """
+    Add a vertical seam to a 3-channel color image at the indices provided 
+    by averaging the pixels values to the left and right of the seam.
+    Code adapted from https://github.com/vivianhylee/seam-carving.
+    """
+    h, w = im.shape[:2]
+    output = np.zeros((h, w + 1, 3))
+    for row in range(h):
+        col = seam_idx[row]
+        for ch in range(3):
+            if col == 0:
+                p = np.mean(im[row, col: col + 2, ch])
+                output[row, col, ch] = im[row, col, ch]
+                output[row, col + 1, ch] = p
+                output[row, col + 1:, ch] = im[row, col:, ch]
+            else:
+                p = np.mean(im[row, col - 1: col + 1, ch])
+                output[row, : col, ch] = im[row, : col, ch]
+                output[row, col, ch] = p
+                output[row, col + 1:, ch] = im[row, col:, ch]
+
+    return output
+
+def add_seam_grayscale(im, seam_idx):
+    """
+    Add a vertical seam to a grayscale image at the indices provided 
+    by averaging the pixels values to the left and right of the seam.
+    """    
+    h, w = im.shape[:2]
+    output = np.zeros((h, w + 1))
+    for row in range(h):
+        col = seam_idx[row]
+        if col == 0:
+            p = np.mean(im[row, col: col + 2])
+            output[row, col] = im[row, col]
+            output[row, col + 1] = p
+            output[row, col + 1:] = im[row, col:]
+        else:
+            p = np.mean(im[row, col - 1: col + 1])
+            output[row, : col] = im[row, : col]
+            output[row, col] = p
+            output[row, col + 1:] = im[row, col:]
+
+    return output
+
+def remove_seam(im, boolmask):
+    h, w = im.shape[:2]
+    boolmask3c = np.stack([boolmask] * 3, axis=2)
+    return im[boolmask3c].reshape((h, w - 1, 3))
+
+def remove_seam_grayscale(im, boolmask):
+    h, w = im.shape[:2]
+    return im[boolmask].reshape((h, w - 1))
+
+def get_minimum_seam(im, mask=None, remove_mask=None):
+    """
+    DP algorithm for finding the seam of minimum energy. Code adapted from 
+    https://karthikkaranth.me/blog/implementing-seam-carving-with-python/
+    """
+    h, w = im.shape[:2]
+    energyfn = forward_energy if USE_FORWARD_ENERGY else backward_energy
+    M = energyfn(im)
+
+    if mask is not None:
+        M[np.where(mask > MASK_THRESHOLD)] = ENERGY_MASK_CONST
+
+    # give removal mask priority over protective mask by using larger negative value
+    if remove_mask is not None:
+        M[np.where(remove_mask > MASK_THRESHOLD)] = -ENERGY_MASK_CONST * 100
+
+    seam_idx, boolmask = compute_shortest_path(M, im, h, w)
+
+    return np.array(seam_idx), boolmask
+
+def compute_shortest_path(M, im, h, w):
+    backtrack = np.zeros_like(M, dtype=np.int_)
+
+
+    # populate DP matrix
+    for i in range(1, h):
+        for j in range(0, w):
+            if j == 0:
+                idx = np.argmin(M[i - 1, j:j + 2])
+                backtrack[i, j] = idx + j
+                min_energy = M[i-1, idx + j]
+            else:
+                idx = np.argmin(M[i - 1, j - 1:j + 2])
+                backtrack[i, j] = idx + j - 1
+                min_energy = M[i - 1, idx + j - 1]
+
+            M[i, j] += min_energy
+
+    # backtrack to find path
+    seam_idx = []
+    boolmask = np.ones((h, w), dtype=np.bool_)
+    j = np.argmin(M[-1])
+    for i in range(h-1, -1, -1):
+        boolmask[i, j] = False
+        seam_idx.append(j)
+        j = backtrack[i, j]
+
+    seam_idx.reverse()
+    return seam_idx, boolmask
+
+########################################
+# MAIN ALGORITHM
+######################################## 
+
+def seams_removal(im, num_remove, mask=None, vis=False, rot=False):
+    for _ in range(num_remove):
+        seam_idx, boolmask = get_minimum_seam(im, mask)
+        if vis:
+            visualize(im, boolmask, rotate=rot)
+        im = remove_seam(im, boolmask)
+        if mask is not None:
+            mask = remove_seam_grayscale(mask, boolmask)
+    return im, mask
+
+
+def seams_insertion(im, num_add, mask=None, vis=False, rot=False):
+    seams_record = []
+    temp_im = im.copy()
+    temp_mask = mask.copy() if mask is not None else None
+
+    for _ in range(num_add):
+        seam_idx, boolmask = get_minimum_seam(temp_im, temp_mask)
+        if vis:
+            visualize(temp_im, boolmask, rotate=rot)
+
+        seams_record.append(seam_idx)
+        temp_im = remove_seam(temp_im, boolmask)
+        if temp_mask is not None:
+            temp_mask = remove_seam_grayscale(temp_mask, boolmask)
+
+    seams_record.reverse()
+
+    for _ in range(num_add):
+        seam = seams_record.pop()
+        im = add_seam(im, seam)
+        if vis:
+            visualize(im, rotate=rot)
+        if mask is not None:
+            mask = add_seam_grayscale(mask, seam)
+
+        # update the remaining seam indices
+        for remaining_seam in seams_record:
+            remaining_seam[np.where(remaining_seam >= seam)] += 2         
+
+    return im, mask
+
+########################################
+# MAIN DRIVER FUNCTIONS
+########################################
+
+def seam_carve(im, dy, dx, mask=None, vis=False):
+    im = im.astype(np.float64)
+    h, w = im.shape[:2]
+    assert h + dy > 0 and w + dx > 0 and dy <= h and dx <= w
+
+    if mask is not None:
+        mask = mask.astype(np.float64)
+
+    output = im
+
+    if dx < 0:
+        output, mask = seams_removal(output, -dx, mask, vis)
+
+    elif dx > 0:
+        output, mask = seams_insertion(output, dx, mask, vis)
+
+    if dy < 0:
+        output = rotate_image(output, True)
+        if mask is not None:
+            mask = rotate_image(mask, True)
+        output, mask = seams_removal(output, -dy, mask, vis, rot=True)
+        output = rotate_image(output, False)
+
+    elif dy > 0:
+        output = rotate_image(output, True)
+        if mask is not None:
+            mask = rotate_image(mask, True)
+        output, mask = seams_insertion(output, dy, mask, vis, rot=True)
+        output = rotate_image(output, False)
+
+    return output
+
+
+def object_removal(im, rmask, mask=None, vis=False, horizontal_removal=False):
+    im = im.astype(np.float64)
+    rmask = rmask.astype(np.float64)
+    if mask is not None:
+        mask = mask.astype(np.float64)
+    output = im
+
+    h, w = im.shape[:2]
+
+    if horizontal_removal:
+        output = rotate_image(output, True)
+        rmask = rotate_image(rmask, True)
+        if mask is not None:
+            mask = rotate_image(mask, True)
+
+    while len(np.where(rmask > MASK_THRESHOLD)[0]) > 0:
+        seam_idx, boolmask = get_minimum_seam(output, mask, rmask)
+        if vis:
+            visualize(output, boolmask, rotate=horizontal_removal)            
+        output = remove_seam(output, boolmask)
+        rmask = remove_seam_grayscale(rmask, boolmask)
+        if mask is not None:
+            mask = remove_seam_grayscale(mask, boolmask)
+
+    num_add = (h if horizontal_removal else w) - output.shape[1]
+    output, mask = seams_insertion(output, num_add, mask, vis, rot=horizontal_removal)
+    if horizontal_removal:
+        output = rotate_image(output, False)
+
+    return output        
+
+
+
+def s_image(im,mask,vs,hs,mode="resize"):
+    im = cv2.cvtColor(im, cv2.COLOR_RGBA2RGB)
+    mask = 255-mask[:,:,3]
+    h, w = im.shape[:2]
+    if SHOULD_DOWNSIZE and w > DOWNSIZE_WIDTH:
+        im = resize(im, width=DOWNSIZE_WIDTH)
+        if mask is not None:
+            mask = resize(mask, width=DOWNSIZE_WIDTH)
+
+    # image resize mode
+    if mode=="resize":
+        dy = hs#reverse
+        dx = vs#reverse
+        assert dy is not None and dx is not None
+        output = seam_carve(im, dy, dx, mask, False)
+        
+
+    # object removal mode
+    elif mode=="remove":
+        assert mask is not None
+        output = object_removal(im, mask, None, False, True)
+        
+    return output
+
+
+##### Inpainting helper code
+
+def run(image, mask):
+    """
+    image: [C, H, W]
+    mask: [1, H, W]
+    return: BGR IMAGE
+    """
+    origin_height, origin_width = image.shape[1:]
+    image = pad_img_to_modulo(image, mod=8)
+    mask = pad_img_to_modulo(mask, mod=8)
+
+    mask = (mask > 0) * 1
+    image = torch.from_numpy(image).unsqueeze(0).to(device)
+    mask = torch.from_numpy(mask).unsqueeze(0).to(device)
+
+    start = time.time()
+    with torch.no_grad():
+        inpainted_image = model(image, mask)
+
+    print(f"process time: {(time.time() - start)*1000}ms")
+    cur_res = inpainted_image[0].permute(1, 2, 0).detach().cpu().numpy()
+    cur_res = cur_res[0:origin_height, 0:origin_width, :]
+    cur_res = np.clip(cur_res * 255, 0, 255).astype("uint8")
+    cur_res = cv2.cvtColor(cur_res, cv2.COLOR_BGR2RGB)
+    return cur_res
+
+
+def get_args_parser():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", default=8080, type=int)
+    parser.add_argument("--device", default="cuda", type=str)
+    parser.add_argument("--debug", action="store_true")
+    return parser.parse_args()
+
+
+def process_inpaint(image, mask):
+    image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+    original_shape = image.shape
+    interpolation = cv2.INTER_CUBIC
+
+    #size_limit: Union[int, str] = request.form.get("sizeLimit", "1080")
+    #if size_limit == "Original":
+    size_limit = max(image.shape)
+    #else:
+    #    size_limit = int(size_limit)
+
+    print(f"Origin image shape: {original_shape}")
+    image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation)
+    print(f"Resized image shape: {image.shape}")
+    image = norm_img(image)
+
+    mask = 255-mask[:,:,3]
+    mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation)
+    mask = norm_img(mask)
+
+    res_np_img = run(image, mask)
+
+    return cv2.cvtColor(res_np_img, cv2.COLOR_BGR2RGB)

src/helper.py

@@ -0,0 +1,87 @@
+import os
+import sys
+
+from urllib.parse import urlparse
+import cv2
+import numpy as np
+import torch
+from torch.hub import download_url_to_file, get_dir
+
+LAMA_MODEL_URL = os.environ.get(
+    "LAMA_MODEL_URL",
+    "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
+)
+
+
+def download_model(url=LAMA_MODEL_URL):
+    parts = urlparse(url)
+    hub_dir = get_dir()
+    model_dir = os.path.join(hub_dir, "checkpoints")
+    if not os.path.isdir(model_dir):
+        os.makedirs(os.path.join(model_dir, "hub", "checkpoints"))
+    filename = os.path.basename(parts.path)
+    cached_file = os.path.join(model_dir, filename)
+    if not os.path.exists(cached_file):
+        sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
+        hash_prefix = None
+        download_url_to_file(url, cached_file, hash_prefix, progress=True)
+    return cached_file
+
+
+def ceil_modulo(x, mod):
+    if x % mod == 0:
+        return x
+    return (x // mod + 1) * mod
+
+
+def numpy_to_bytes(image_numpy: np.ndarray) -> bytes:
+    data = cv2.imencode(".jpg", image_numpy)[1]
+    image_bytes = data.tobytes()
+    return image_bytes
+
+
+def load_img(img_bytes, gray: bool = False):
+    nparr = np.frombuffer(img_bytes, np.uint8)
+    if gray:
+        np_img = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE)
+    else:
+        np_img = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
+        if len(np_img.shape) == 3 and np_img.shape[2] == 4:
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGRA2RGB)
+        else:
+            np_img = cv2.cvtColor(np_img, cv2.COLOR_BGR2RGB)
+
+    return np_img
+
+
+def norm_img(np_img):
+    if len(np_img.shape) == 2:
+        np_img = np_img[:, :, np.newaxis]
+    np_img = np.transpose(np_img, (2, 0, 1))
+    np_img = np_img.astype("float32") / 255
+    return np_img
+
+
+def resize_max_size(
+    np_img, size_limit: int, interpolation=cv2.INTER_CUBIC
+) -> np.ndarray:
+    # Resize image's longer size to size_limit if longer size larger than size_limit
+    h, w = np_img.shape[:2]
+    if max(h, w) > size_limit:
+        ratio = size_limit / max(h, w)
+        new_w = int(w * ratio + 0.5)
+        new_h = int(h * ratio + 0.5)
+        return cv2.resize(np_img, dsize=(new_w, new_h), interpolation=interpolation)
+    else:
+        return np_img
+
+
+def pad_img_to_modulo(img, mod):
+    channels, height, width = img.shape
+    out_height = ceil_modulo(height, mod)
+    out_width = ceil_modulo(width, mod)
+    return np.pad(
+        img,
+        ((0, 0), (0, out_height - height), (0, out_width - width)),
+        mode="symmetric",
+    )

src/st_style.py

@@ -0,0 +1,42 @@
+button_style = """
+<style>
+div.stButton > button:first-child {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+div.stButton > button:hover {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+div.stButton > button:active {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+div.stButton > button:focus {
+    background-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+.css-1cpxqw2:focus:not(:active) {
+    background-color: rgb(255, 75, 75);
+    border-color: rgb(255, 75, 75);
+    color: rgb(255, 255, 255);
+}
+"""
+
+style = """
+<style>
+#MainMenu {
+    visibility: hidden;
+}
+footer {
+    visibility: hidden;
+}
+header {
+    visibility: hidden;
+}
+</style>
+"""
+
+
+def apply_prod_style(st):
+    return st.markdown(style, unsafe_allow_html=True)