Spaces:
Running
on
T4
Running
on
T4
File size: 4,770 Bytes
3fdea04 d5eecf2 3fdea04 e670614 3fdea04 e670614 3fdea04 e670614 0801672 3fdea04 0801672 3fdea04 d0dc890 3fdea04 241ac77 3fdea04 241ac77 3fdea04 241ac77 3fdea04 d5eecf2 3fdea04 d0dc890 3fdea04 e670614 c98ba13 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
# 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
def run_streamlit_ui():
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")
else:
st.info("Please upload an image to begin.")
return
# Resize the image while maintaining aspect ratio
max_size = 2000
img_width, img_height = img_input.size
if img_width > max_size or img_height > max_size:
if img_width > img_height:
new_width = max_size
new_height = int((max_size / img_width) * img_height)
else:
new_height = max_size
new_width = int((max_size / img_height) * img_width)
img_input = img_input.resize((new_width, new_height))
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.astype(np.uint8)).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
if st.button('Submit'):
with st.spinner("AI is doing the magic!"):
output = process_inpaint(np.array(img_input), np.array(im))
img_output = Image.fromarray(output).convert("RGB")
st.write("AI has finished the job!")
st.image(img_output)
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.")
# Initialize Streamlit app (HF Spaces runs via `streamlit run app.py`)
run_streamlit_ui()
|