Spaces:
Runtime error
Runtime error
| import os | |
| from PIL import Image | |
| import streamlit as st | |
| from tensorflow.keras import utils, backend | |
| from enhance_me import MIRNet, ZeroDCE | |
| def get_mirnet_object() -> MIRNet: | |
| utils.get_file( | |
| "weights_lol_128.h5", | |
| "https://github.com/soumik12345/enhance-me/releases/download/v0.2/weights_lol_128.h5", | |
| cache_dir=".", | |
| cache_subdir="weights", | |
| ) | |
| mirnet = MIRNet() | |
| mirnet.build_model() | |
| mirnet.load_weights("./weights/weights_lol_128.h5") | |
| return mirnet | |
| def get_zero_dce_object(model_alias: str) -> ZeroDCE: | |
| utils.get_file( | |
| f"{model_alias}.h5", | |
| f"https://github.com/soumik12345/enhance-me/releases/download/v0.4/{model_alias}.h5", | |
| cache_dir=".", | |
| cache_subdir="weights", | |
| ) | |
| dce = ZeroDCE() | |
| dce.load_weights(os.path.join("./weights", f"{model_alias}.h5")) | |
| return dce | |
| def main(): | |
| st.markdown("# Enhance Me") | |
| st.markdown("Made with :heart: by [geekyRakshit](http://github.com/soumik12345)") | |
| application = st.sidebar.selectbox( | |
| "Please select the application:", ("", "Low-light enhancement") | |
| ) | |
| if application != "": | |
| if application == "Low-light enhancement": | |
| uploaded_file = st.sidebar.file_uploader("Select your image:") | |
| if uploaded_file is not None: | |
| original_image = Image.open(uploaded_file) | |
| st.image(original_image, caption="original image") | |
| model_option = st.sidebar.selectbox( | |
| "Please select the model:", | |
| ( | |
| "", | |
| "MIRNet", | |
| "Zero-DCE (dce_weights_lol_128)", | |
| "Zero-DCE (dce_weights_lol_128_resize)", | |
| "Zero-DCE (dce_weights_lol_256)", | |
| "Zero-DCE (dce_weights_lol_256_resize)", | |
| "Zero-DCE (dce_weights_unpaired_128)", | |
| "Zero-DCE (dce_weights_unpaired_128_resize)", | |
| "Zero-DCE (dce_weights_unpaired_256)", | |
| "Zero-DCE (dce_weights_unpaired_256_resize)" | |
| ), | |
| ) | |
| if model_option != "": | |
| if model_option == "MIRNet": | |
| st.sidebar.info("Loading MIRNet...") | |
| mirnet = get_mirnet_object() | |
| st.sidebar.info("Done!") | |
| st.sidebar.info("Processing Image...") | |
| enhanced_image = mirnet.infer(original_image) | |
| st.sidebar.info("Done!") | |
| st.image(enhanced_image, caption="enhanced image") | |
| elif "Zero-DCE" in model_option: | |
| model_alias = model_option[model_option.find("(") + 1: model_option.find(")")] | |
| st.sidebar.info("Loading Zero-DCE...") | |
| zero_dce = get_zero_dce_object(model_alias) | |
| st.sidebar.info("Done!") | |
| enhanced_image = zero_dce.infer(original_image) | |
| st.sidebar.info("Done!") | |
| st.image(enhanced_image, caption="enhanced image") | |
| backend.clear_session() | |
| if __name__ == "__main__": | |
| main() | |