Spaces:
Sleeping
Sleeping
| 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() |