|
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() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
results_view(name, path) |
|
if uploaded_file is not None: |
|
name, path = save_uploadedfile(uploaded_file) |
|
results_view(name, path) |
|
|
|
if menu == 'About': |
|
about() |
