| from libs import * |
| from utils_func import create_dir, main_processing |
|
|
|
|
|
|
| create_dir("tempDir") |
|
|
|
|
| def load_image(image_file): |
| img = Image.open(image_file) |
| return img |
|
|
|
|
| def streamlit_app(): |
| detection_model_path = "weight_files/clothes_detection_model.pt" |
| background_model_path = "weight_files/model.h5" |
| save_path = "" |
| image_file = None |
| st.title("""WELCOME TO MY APP""") |
| st.subheader("""FOR BACKGROUND REMOVAL AND CHANGE!""") |
| col1 = None |
| col2 = None |
| final_img = None |
| with st.spinner("[UPLOAD] Image uploading"): |
| try: |
| image_file = st.file_uploader('[UPLOAD] Please upload your image:', type=["png", "jpg", "jpeg"]) |
| time.sleep(1) |
| except: |
| print("[ERROR] Sorry, something went wrong!") |
| pass |
| |
|
|
| if image_file is not None: |
| st.success("Load image successfully!...") |
| image = load_image(image_file) |
| |
| col1, col2, col3 = st.columns(3) |
| with col1: |
| st.image(image, caption="Image before processing") |
| save_path = "tempDir/"+ image_file.name |
| image.save(save_path) |
|
|
|
|
| image_path, details = save_path, image_file |
|
|
| if details is not None: |
| with col2: |
| with st.spinner("[PROCESSING] Image processing"): |
| final_img_path = main_processing(col1, col2, col3, sport_bg_path=stadium_sport_bg_path, swim_bg_path=beach_swim_bg_path, |
| office_bg_path=office_bg_path, img_path=image_path, name=details.name, |
| detection_model_path=detection_model_path, |
| background_model_path=background_model_path) |
| time.sleep(1) |
|
|
| with col1: |
| if final_img_path is not None: |
| final_img = load_image(final_img_path) |
| st.image(final_img, caption="Image after processing") |
| st.balloons() |
| with col2: |
| with open(final_img_path, "rb") as file: |
| st.write('\n') |
| st.write('\n') |
| st.write('\n') |
| st.write('\n') |
| st.write('\n') |
|
|
| file_name = save_path.split("/")[-1].split(".")[-2] +"_from_abc" + ".png" |
|
|
| if st.download_button( |
| label="Download postprocessing image", |
| data=file, |
| file_name= file_name, |
| mime="image/png" |
| ): |
| st.success('[DOWNLOAD] Download sucessfully!') |
|
|
|
|
|
|
| if __name__ == '__main__': |
| np.random.seed(42) |
| tf.random.set_seed(42) |
|
|
| bg_path = "" |
| background_model_path = "weight_files/model.h5" |
| detection_model_path = "weight_files/clothes_detection_model.pt" |
|
|
| stadium_sport_bg_path = "backgrounds/camnou_stadium.jpg" |
| beach_swim_bg_path = "backgrounds/beach.jpg" |
| office_bg_path = "backgrounds/office-bg.jpg" |
|
|
| image_path = None |
|
|
| streamlit_app() |
|
|
