Spaces:
Running
Running
File size: 3,984 Bytes
e0ac091 |
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 |
"""Application to demo inpainting using streamlit.
Run using: streamlit run 10_03_image_inpaint_streamlit.py
"""
import streamlit as st
import pathlib
from streamlit_drawable_canvas import st_canvas
import cv2
import numpy as np
import io
import base64
from PIL import Image
STREAMLIT_STATIC_PATH = pathlib.Path(st.__path__[0]) / 'static'
# We create a downloads directory within the streamlit static asset directory
# and we write output files to it.
DOWNLOADS_PATH = (STREAMLIT_STATIC_PATH / "downloads")
if not DOWNLOADS_PATH.is_dir():
DOWNLOADS_PATH.mkdir()
def get_image_download_link(img, filename, text):
"""Generates a link to download a particular image file."""
buffered = io.BytesIO()
img.save(buffered, format='JPEG')
img_str = base64.b64encode(buffered.getvalue()).decode()
href = f'<a href="data:file/txt;base64,{img_str}" download="{filename}">{text}</a>'
return href
# Set title.
st.sidebar.title('Image Inpaint')
# Specify canvas parameters in application
uploaded_file = st.sidebar.file_uploader("Upload Image to restore:", type=["png", "jpg"])
image = None
res = None
if uploaded_file is not None:
# Convert the file to an opencv image.
file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
image = cv2.imdecode(file_bytes, 1)
stroke_width = st.sidebar.slider("Stroke width: ", 1, 25, 5)
h, w = image.shape[:2]
if w > 800:
h_, w_ = int(h * 800 / w), 800
else:
h_, w_ = h, w
# Create a canvas component.
canvas_result = st_canvas(
fill_color='gray',
stroke_width=stroke_width,
stroke_color='green',
background_image=Image.open(uploaded_file).resize((h_, w_)),
update_streamlit=True,
height=h_,
width=w_,
drawing_mode='freedraw',
key="canvas",
)
stroke = canvas_result.image_data
if stroke is not None:
if st.sidebar.checkbox('show mask'):
st.image(stroke)
mask = cv2.split(stroke)[3]
mask = np.uint8(mask)
mask = cv2.resize(mask, (w, h))
st.sidebar.caption('Happy with the selection?')
option = st.sidebar.selectbox('Mode', ['None', 'Telea', 'NS', 'Compare both'])
if option == 'Telea':
st.subheader('Result of Telea')
res = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)[:, :, ::-1]
st.image(res)
elif option == 'Compare both':
col1, col2 = st.columns(2)
res1 = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)[:, :, ::-1]
res2 = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_NS)[:, :, ::-1]
with col1:
st.subheader('Result of Telea')
st.image(res1)
with col2:
st.subheader('Result of NS')
st.image(res2)
if res1 is not None:
# Display link.
result1 = Image.fromarray(res1)
st.sidebar.markdown(
get_image_download_link(result1, 'telea.png', 'Download Output of Telea'),
unsafe_allow_html=True)
if res2 is not None:
# Display link.
result2 = Image.fromarray(res2)
st.sidebar.markdown(
get_image_download_link(result2, 'ns.png', 'Download Output of NS'),
unsafe_allow_html=True)
elif option == 'NS':
st.subheader('Result of NS')
res = cv2.inpaint(src=image, inpaintMask=mask, inpaintRadius=3, flags=cv2.INPAINT_NS)[:, :, ::-1]
st.image(res)
else:
pass
if res is not None:
# Display link.
result = Image.fromarray(res)
st.sidebar.markdown(
get_image_download_link(result, 'output.png', 'Download Output'),
unsafe_allow_html=True)
|