File size: 3,660 Bytes
7439e48 |
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 |
import random
from PIL import Image
import os
import shutil
import streamlit as st
from about import about
st.set_page_config(
page_title=None,
page_icon=None,
layout="wide",
initial_sidebar_state="auto",
menu_items=None,
)
st.title("Image Enhancer")
if os.path.exists("results"):
shutil.rmtree(os.path.join("results"))
if os.path.exists("tempDir"):
shutil.rmtree(os.path.join("tempDir"))
def create_dir(dirname: str):
if not os.path.exists(dirname):
os.makedirs(dirname, exist_ok=True)
create_dir("results/cmp")
create_dir("results/cropped_faces")
create_dir("results/restored_faces")
create_dir("results/restored_imgs")
create_dir("tempDir")
def save_uploadedfile(uploadedfile):
file_extension = os.path.splitext(uploadedfile.name)[-1].lstrip(".")
with open(os.path.join("tempDir", f"uploaded_image.{file_extension}"), "wb") as f:
f.write(uploadedfile.getbuffer())
name, path = f'uploaded_image.{file_extension}', os.path.join("tempDir", f"uploaded_image.{file_extension}")
return name, path
def get_random_sample_image():
sample_images = os.listdir(os.path.join('sample_images'))
random_image = random.choice(sample_images)
return random_image, f"{os.getcwd()}/sample_images/{random_image}"
def results_view(name, path):
with st.spinner("Please wait while we process your image.."):
os.system(f"python inference_gfpgan.py -i {path} -o results -v {version} -s 2")
with st.expander("Results", expanded=True):
col_1_1, col_2_2 = st.columns(2)
with col_1_1:
st.write("Sample Image")
st.image(path)
with col_2_2:
st.write("Processed Image")
st.image(Image.open(os.path.join("results", "restored_imgs", name)))
with st.expander("Comparative Results", expanded=True):
files = os.listdir(os.path.join("results", "cmp"))
for f in files:
st.write(f)
st.image(Image.open(os.path.join("results", "cmp", f)))
options = ["Face Restore", "About"]
models = ['1.3', '1.4', "RestoreFormer"]
st.sidebar.image("assets/gfpgan_logo.png")
menu = st.sidebar.selectbox("Select an Option", options)
if menu == "Face Restore":
col1, col2 = st.columns([1, 0.3])
with st.sidebar:
version = st.selectbox("Select Version", models)
with col1:
uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
with col2:
st.markdown("###")
st.markdown("###")
sample = st.button("Use Sample Image")
if sample:
name, path = get_random_sample_image()
# with st.spinner("Please wait while we process your image.."):
# os.system(f"python inference_gfpgan.py -i {path} -o results -v {version} -s 2")
# with st.expander("Results", expanded=True):
# col_1_1, col_2_2 = st.columns(2)
# with col_1_1:
# st.write("Sample Image")
# st.image(path)
# with col_2_2:
# st.write("Processed Image")
# st.image(Image.open(os.path.join("results", "restored_imgs", name)))
# with st.expander("Comparative Results", expanded=True):
# files = os.listdir(os.path.join("results", "cmp"))
# for f in files:
# st.write(f)
# st.image(Image.open(os.path.join("results", "cmp", f)))
results_view(name, path)
if uploaded_file is not None:
name, path = save_uploadedfile(uploaded_file)
results_view(name, path)
if menu == 'About':
about() |