Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,53 +1,78 @@
-import tempfile
 import streamlit as st
+import numpy as np
+from PIL import Image, ImageOps
 import tensorflow as tf
 from tensorflow import keras
-import numpy as np
-from PIL import Image
 import io
-from huggingface_hub import hf_hub_download
-import os
-
-# Set Streamlit page
-st.set_page_config(page_title="NaxiLowLight Enhancer", layout="centered")
-st.title("🌙 NaxiLowLight - Low-Light Image Enhancer")
 
-# ✅ Load model using temporary directory (no caching to /.cache!)
+# ----------------------------
+# Load the model
+# ----------------------------
 @st.cache_resource
-def load_model_from_huggingface_temp():
-    repo_id = "NEXAS/low_light_enhance"
-    filename = "model.h5"
-
-    with tempfile.TemporaryDirectory() as temp_dir:
-        model_path = hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=temp_dir)
-        model = keras.models.load_model(model_path, compile=False)
+def load_model():
+    model = keras.models.load_model(
+        "src/naxi_lowlight.keras",
+        custom_objects={
+            "charbonnier_loss": lambda y_true, y_pred: tf.reduce_mean(tf.sqrt(tf.square(y_true - y_pred) + 1e-6)),
+            "psnr_metric": lambda y_true, y_pred: tf.image.psnr(y_pred, y_true, max_val=1.0)
+        }
+    )
     return model
 
-model = load_model_from_huggingface_temp()
+model = load_model()
 
-# 🔧 Enhancement function
-def enhance_image(image: Image.Image) -> Image.Image:
-    image_resized = image.resize((256, 256))
-    img_array = keras.utils.img_to_array(image_resized) / 255.0
-    img_array = np.expand_dims(img_array, axis=0)
-    output = model.predict(img_array)[0]
+# ----------------------------
+# Inference function
+# ----------------------------
+def enhance_image_pil(pil_img):
+    image = keras.utils.img_to_array(pil_img).astype("float32") / 255.0
+    image = np.expand_dims(image, axis=0)
+    output = model.predict(image)[0]
     output = np.clip(output * 255.0, 0, 255).astype(np.uint8)
     return Image.fromarray(output)
 
-# 📤 Upload & Enhance UI
-uploaded_file = st.file_uploader("📤 Upload a low-light image", type=["jpg", "jpeg", "png"])
+# ----------------------------
+# Streamlit UI
+# ----------------------------
+st.set_page_config(page_title="NaxiLowLight Enhancement", layout="centered")
+st.title("🌙 NaxiLowLight: Low-Light Image Enhancer")
+
+st.write("Upload a low-light image to enhance it using a deep learning model.")
 
-if uploaded_file:
+uploaded_file = st.file_uploader("Choose a low-light image", type=["jpg", "jpeg", "png"])
+
+if uploaded_file is not None:
+    # Load image
     image = Image.open(uploaded_file).convert("RGB")
-    st.image(image, caption="🖼️ Original Image", use_column_width=True)
 
-    if st.button("✨ Enhance Image"):
-        with st.spinner("Enhancing image..."):
-            enhanced = enhance_image(image)
+    # Enhance
+    st.write("✨ Enhancing image...")
+    enhanced_image = enhance_image_pil(image)
+    autocontrast_image = ImageOps.autocontrast(image)
+
+    # Show results
+    st.write("📷 **Comparison:**")
+    col1, col2, col3 = st.columns(3)
+
+    with col1:
+        st.image(image, caption="Original", use_column_width=True)
+
+    with col2:
+        st.image(autocontrast_image, caption="Autocontrast", use_column_width=True)
+
+    with col3:
+        st.image(enhanced_image, caption="NaxiLowLight Enhanced", use_column_width=True)
 
-        st.image(enhanced, caption="🚀 Enhanced Image", use_column_width=True)
+    # Download
+    buf = io.BytesIO()
+    enhanced_image.save(buf, format="PNG")
+    byte_im = buf.getvalue()
 
-        # Download
-        img_bytes = io.BytesIO()
-        enhanced.save(img_bytes, format="PNG")
-        st.download_button("💾 Download Enhanced Image", data=img_bytes.getvalue(), file_name="enhanced.png", mime="image/png")
+    st.download_button(
+        label="⬇️ Download Enhanced Image",
+        data=byte_im,
+        file_name="enhanced_image.png",
+        mime="image/png"
+    )
+else:
+    st.info("Upload a PNG or JPG image to get started.")