ReflectionEraser/ReflectionEraserApp

1import streamlit as st2import subprocess3from PIL import Image4import tempfile5import os6 7# Streamlit app title8st.title("Reflection Removal App")9 10# File upload section11uploaded_file = st.file_uploader("Choose a file", type=["jpg", "jpeg", "png"])12 13# Display uploaded image if available14if uploaded_file is not None:15    # Save the uploaded file to a temporary directory16    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:17        tmp_file.write(uploaded_file.getvalue())18        tmp_file_path = tmp_file.name19 20    #writing to text file21    #22    #########################23    file = open("uploaded_img.txt", "w")24 25    # Write data to the file26    file.write(tmp_file.name)27    file.write("This is a second line.\n")28    29    # Close the file30    file.close()31 32 33    34    # Display the uploaded image35    image = Image.open(tmp_file_path)36    st.image(image, caption='Uploaded Image', use_column_width=True)37    38    # Define constants or variables for model execution39    MODEL_NAME = 'dsrnet_s'40    MODEL_PATH = 'dsrnet_model_sirs'41    DATASET = 'sirs_dataset'42    NAME = 'dsrnet_s_test'43    BASE_DIR = './real45'44    WEIGHT_PATH = './weights/dsrnet_s_epoch14.pt'45    46    # Function to run the command47    def run_command(command):48        try:49            process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)50            stdout, stderr = process.communicate()51            return stdout.decode('utf-8'), stderr.decode('utf-8')52        except Exception as e:53            return "", str(e)54    55    # Command to execute56    command = f"python ./DSRNet/test_sirs.py --inet {MODEL_NAME} --model {MODEL_PATH} --dataset {DATASET} --name {NAME} --hyper --if_align --base_dir {BASE_DIR} --resume --weight_path {WEIGHT_PATH}"57    58    # Display running command59    st.write(f"Running command: {command}")60    61    # Execute the command62    stdout, stderr = run_command(command)63    64    # Display command output65    st.write("Command Output:")66    if stdout:67        st.code(stdout)68    if stderr:69        st.error(f"Command Error:\n{stderr}")70 71    # Search for image ending with '_l' in subdirectories of 'checkpoints'72    checkpoints_dir = './test'  # Update this path according to your project structure73 74    # Function to recursively find image ending with '_l'75 76    # Find the image file77    desired_image = './test/dsrnet_s_l.png'78 79    # Display the found image if available80    if desired_image is not None:81        st.write("Found Image:")82        image_l = Image.open(desired_image)83        st.image(image_l, caption='Image ending with _l', use_column_width=True)84    else:85        st.write("Image ending with _l not found in checkpoints directory.")86 87