Inpaint_Me / app_inpaint.py
Huertas97's picture
Remove prints
1f7249f
raw
history blame
6.45 kB
import streamlit as st
import time
import easyocr
import math
from pathlib import Path
from PIL import Image, ImageDraw
import PIL
import io
import os
import cv2
import numpy as np
import shutil
import base64
import logging
@st.cache(show_spinner=False, allow_output_mutation=True, suppress_st_warning=True)
def load_models():
#specify shortform of language you want to extract,
# I am using Spanish(es) and English(en) here by list of language ids
reader = easyocr.Reader(['en'],)
return reader
reader = load_models()
def midpoint(x1, y1, x2, y2):
x_mid = int((x1 + x2)/2)
y_mid = int((y1 + y2)/2)
return (x_mid, y_mid)
def inpaint_text(img, text_coordinates):
# read image
# img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# generate (word, box) tuples
mask = np.zeros(img.shape[:2], dtype="uint8")
for box in text_coordinates:
x0, y0 = box[0]
x1, y1 = box[1]
x2, y2 = box[2]
x3, y3 = box[3]
x_mid0, y_mid0 = midpoint(x1, y1, x2, y2)
x_mid1, y_mi1 = midpoint(x0, y0, x3, y3)
thickness = int(math.sqrt( (x2 - x1)**2 + (y2 - y1)**2 ))
cv2.line(mask, (x_mid0, y_mid0), (x_mid1, y_mi1), 255,
thickness)
img = cv2.inpaint(img, mask, 7, cv2.INPAINT_NS)
return(img)
def file_selector(folder_path='.'):
filenames = os.listdir(folder_path)
selected_filename = st.selectbox('Select a file', filenames)
return os.path.join(folder_path, selected_filename), selected_filename
st.set_page_config(
page_title="Inpaint Me",
page_icon=":art:",
layout="wide",
initial_sidebar_state="expanded",
menu_items={
'Get Help': 'https://www.extremelycoolapp.com/help',
'Report a bug': "https://www.extremelycoolapp.com/bug",
'About': "# This is a header. This is an *extremely* cool app!"
}
)
st.markdown(
"""
<style>
.logo-img {
margin-top: auto;
margin-left: 30%;
width: 50%;
}
.logo-img-2 {
margin-top: auto;
margin-left: 20%;
width: 35%;
}
</style>
""",
unsafe_allow_html=True
)
LOGO_IMAGE = "inpaint_me_logo.png"
col1, col2= st.columns([2, 2])
with col1:
# st.image('./aida_logo.png')
st.markdown(
f"""
<img class="logo-img" src="data:image/png;base64,{base64.b64encode(open(LOGO_IMAGE, "rb").read()).decode()}">
""",
unsafe_allow_html=True
)
with col2:
# st.image('./aida_logo.png')
st.markdown(
f"""
<img class="logo-img-2" src="data:image/png;base64,{base64.b64encode(open("aida_logo.png", "rb").read()).decode()}">
""",
unsafe_allow_html=True
)
st.header("")
with st.expander("Project Description", expanded=False):
st.write("""
Developed in Applied Intelligence and Data Analysis ([AI+DA](http://aida.etsisi.upm.es/)) group at Polytech University of Madrid (UPM).
To rule out the possibility of text misleading image Deep Learning models (e.g., CNNs) it is useful to remove text from images. Hence,
this tool uses [EasyOCR](https://github.com/JaidedAI/EasyOCR) and [OpenCV](https://pypi.org/project/opencv-python/) for detecting texts and inpainting them. Currently, only `JPG` files are supported. This tools has been tested on memes, feel free to try some examples or upload your own images.
""")
file_example_path = None
if st.checkbox('Select a example'):
folder_path = './Examples/'
# if st.checkbox('Change directory'):
# folder_path = st.text_input('Enter folder path', '.')
file_example_path, example_file_name = file_selector(folder_path=folder_path)
st.write('You selected `%s`' % file_example_path)
uploaded_file = st.file_uploader(label="Upload image",
type=["jpg", "jpeg"],
accept_multiple_files=False,
key=None,
help=None,
on_change=None,
args=None,
kwargs=None,
)
col1, col2, col3 = st.columns([2, 0.5, 2])
if file_example_path and not uploaded_file:
with col1:
st.subheader("Original")
# st.write(f"./Examples_inpainted/{example_file_name.strip(".jpg")}_inpainted.jpeg")
img = Image.open( file_example_path )
st.image(img, caption=None, width=None, use_column_width=None, clamp=False, channels="RGB", output_format="auto")
with col3:
st.subheader("Inpainted")
with st.spinner('Wait for it...'):
time.sleep(1)
example_file_name = example_file_name.strip(".jpg")
inpaint_image = f"./Examples_inpainted/{example_file_name}_inpaint.jpeg"
# img_array = np.array(Image.open( file_example_path ))
# # detect text
# bounds = reader.readtext(img_array, detail=1) #detail=1 # [(coordinates, detected text, confidence threshold)]
# text_coordinates = [ bound[0] for bound in bounds]
# # inpaint text coordinates
# inpaint_image = inpaint_text(img_array, text_coordinates)
st.image(inpaint_image, caption=None, width=None, use_column_width=None, clamp=False, channels="RGB", output_format="auto")
if uploaded_file:
with col1:
st.subheader("Original")
st.image(uploaded_file, caption=None, width=None, use_column_width=None, clamp=False, channels="RGB", output_format="auto")
with col3:
st.subheader("Inpainted")
with st.spinner('Wait for it...'):
# Transform loaded file to bytes
bytes_data = uploaded_file.getvalue()
# bytes to numpy array
img_array = np.array(Image.open(io.BytesIO(bytes_data)))
# detect text
bounds = reader.readtext(img_array, detail=1) #detail=1 # [(coordinates, detected text, confidence threshold)]
text_coordinates = [ bound[0] for bound in bounds]
# inpaint text coordinates
inpaint_image = inpaint_text(img_array, text_coordinates)
st.image(inpaint_image, caption=None, width=None, use_column_width=None, clamp=False, channels="RGB", output_format="auto")