Upload 2 files

requirements.txt

@@ -1,3 +1,7 @@
-altair
-pandas
-streamlit
+streamlit
+numpy
+opencv-python
+torch
+torchvision
+Pillow
+scipy

seathruapp.py

@@ -0,0 +1,232 @@
+import streamlit as st
+import numpy as np
+import cv2
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import io
+
+# --- Enhancement Functions ---
+def enhance_channel(I_c, z, veil, backscatter, recover, attenuation):
+    """
+    Enhance a single color channel using provided parameters.
+    
+    Args:
+        I_c: Observed image channel (2D array)
+        z: Depth map (2D array)
+        veil: Veiling light (B_c^∞)
+        backscatter: Backscatter coefficient (β_c^B)
+        recover: Recovery factor
+        attenuation: Attenuation coefficient (β_c^D)
+    
+    Returns:
+        Enhanced channel (2D array)
+    """
+    B_c = veil * (1 - np.exp(-backscatter * z))
+    B_c = np.clip(B_c, 0, 1)
+    D_c = np.maximum(I_c - B_c, 0)
+    J_c = D_c * np.exp(attenuation * z) * recover
+    J_c = np.clip(J_c, 0, 10)
+    return J_c
+
+def gray_world_white_balance(image):
+    """
+    Apply Gray World white balancing to normalize colors.
+    
+    Args:
+        image: RGB image (3D array)
+    
+    Returns:
+        White-balanced image
+    """
+    avg_colors = np.mean(image, axis=(0, 1))
+    avg_global = np.mean(avg_colors)
+    if avg_global == 0:
+        return image
+    scaling = avg_global / avg_colors
+    balanced_image = image * scaling[None, None, :]
+    return np.clip(balanced_image, 0, 1)
+
+# --- Depth Estimation with MiDaS ---
+def preprocess_image_for_midas(img):
+    """
+    Preprocess an image for MiDaS depth estimation.
+    Convert to RGB and remove alpha channel if present.
+    
+    Args:
+        img: Input image (numpy array, RGB or RGBA)
+    
+    Returns:
+        RGB image (numpy array)
+    """
+    if len(img.shape) == 3 and img.shape[2] == 4:
+        # If image has 4 channels (RGBA), remove alpha channel
+        img = img[:, :, :3]  # Keep only RGB channels
+    return img
+
+def estimate_depth(frame, depth_model, transform, device, target_shape):
+    """
+    Estimate depth map for an image using MiDaS.
+    
+    Args:
+        frame: Input image (RGB, uint8)
+        depth_model: MiDaS model
+        transform: Preprocessing transform for MiDaS
+        device: Device to run the model on (CPU/GPU)
+        target_shape: Desired output shape (height, width)
+    
+    Returns:
+        Depth map (normalized to [0, 1])
+    """
+    frame_rgb = preprocess_image_for_midas(frame)
+    input_tensor = transform(frame_rgb).to(device)
+
+    with torch.no_grad():
+        depth = depth_model(input_tensor)
+        depth = depth.squeeze().cpu().numpy()
+
+    depth_min, depth_max = depth.min(), depth.max()
+    if depth_max - depth_min > 0:
+        depth = (depth - depth_min) / (depth_max - depth_min)
+    else:
+        depth = np.zeros_like(depth)
+
+    depth = cv2.resize(depth, (target_shape[1], target_shape[0]), interpolation=cv2.INTER_LINEAR)
+    return depth
+
+def load_depth_map(depth_image, target_shape):
+    """
+    Load and preprocess a depth map from an uploaded image.
+    
+    Args:
+        depth_image: PIL Image object
+        target_shape: Desired output shape (height, width)
+    
+    Returns:
+        Depth map (normalized to [0, 1])
+    """
+    depth = np.array(depth_image.convert('L')).astype(np.float32) / 255.0
+    depth = cv2.resize(depth, (target_shape[1], target_shape[0]), interpolation=cv2.INTER_LINEAR)
+    return depth
+
+# --- Streamlit App ---
+def main():
+    st.title("Underwater Image Enhancement with Adjustable Parameters")
+    st.write("Upload an underwater image to begin, then select a mode and adjust the enhancement parameters.")
+
+    # Image uploader (required first)
+    image_file = st.file_uploader("Upload an underwater image", type=["png", "jpg", "jpeg"])
+
+    if image_file is not None:
+        # Load and display the uploaded image
+        image = np.array(Image.open(image_file))
+        h, w = image.shape[:2]
+        st.image(image, caption=f"Original Image ({w}x{h})", use_column_width=True)
+
+        # Convert image for processing (ensure 3 channels)
+        image_float = image.astype(np.float32) / 255.0
+        if len(image_float.shape) == 3 and image_float.shape[2] == 4:
+            image_float = image_float[:, :, :3]  # Remove alpha channel
+
+        # Mode selection (appears after image upload)
+        mode = st.selectbox("Select Mode", ["Map", "Predict", "Hybrid"])
+
+        # Depth image uploader (only for Map or Hybrid mode, after image upload)
+        depth_file = None
+        if mode in ["Map", "Hybrid"]:
+            depth_file = st.file_uploader("Upload the depth map image", type=["png", "jpg", "jpeg"])
+
+        # Initialize MiDaS if needed (Predict or Hybrid mode)
+        depth_model = None
+        transform = None
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if mode in ["Predict", "Hybrid"]:
+            st.write(f"Using device for MiDaS: {device}")
+            with st.spinner("Loading MiDaS model..."):
+                depth_model = torch.hub.load("intel-isl/MiDaS", "DPT_Hybrid", pretrained=True)
+                depth_model = depth_model.to(device).eval()
+                transform = transforms.Compose([
+                    transforms.ToTensor(),
+                    transforms.Resize((384, 384)),
+                    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
+                ])
+
+        # Load or predict depth
+        depth = None
+        if mode == "Map":
+            if depth_file is None:
+                st.error("Please upload a depth map for Map mode.")
+                return
+            depth = load_depth_map(Image.open(depth_file), (h, w))
+        elif mode == "Predict":
+            with st.spinner("Estimating depth with MiDaS..."):
+                depth = estimate_depth(image, depth_model, transform, device, (h, w))
+        elif mode == "Hybrid":
+            if depth_file is None:
+                st.error("Please upload a depth map for Hybrid mode.")
+                return
+            with st.spinner("Estimating depth with MiDaS..."):
+                depth_pred = estimate_depth(image, depth_model, transform, device, (h, w))
+            depth_map = load_depth_map(Image.open(depth_file), (h, w))
+            depth = (depth_map + depth_pred) / 2.0
+
+        # Depth normalization sliders
+        st.subheader("Depth Normalization")
+        depth_min = st.slider("Depth Min", 0.0, 5.0, 0.0, step=0.1)
+        depth_max = st.slider("Depth Max", 0.0, 5.0, 1.0, step=0.1)
+        if depth_max <= depth_min:
+            st.error("Depth Max must be greater than Depth Min")
+            return
+
+        depth_range = depth_max - depth_min
+        depth = depth_min + (depth * depth_range)
+
+        # Parameter sliders for each channel
+        st.subheader("Enhancement Parameters")
+        params = {}
+        channels = ["Red", "Green", "Blue"]
+        for c in range(3):
+            st.write(f"**{channels[c]} Channel**")
+            with st.expander(f"Adjust {channels[c]} Parameters"):
+                params[c] = {
+                    "veil": st.slider(f"Veil ({channels[c]})", 0.0, 1.0, 0.3442605477720724 if c == 0 else 0.36920450457864046 if c == 1 else 0.46370720994475223, step=0.01),
+                    "backscatter": st.slider(f"Backscatter ({channels[c]})", 0.0, 5.0, 0.6980786267220486 if c == 0 else 4.901524448971207 if c == 1 else 4.567895834039181, step=0.01),
+                    "recover": st.slider(f"Recover ({channels[c]})", 0.0, 2.0, 0.9999999999999997 if c == 0 else 0.9999999999999999 if c == 1 else 0.9999999999999998, step=0.01),
+                    "attenuation": st.slider(f"Attenuation ({channels[c]})", 0.0, 1.0, 0.5000000000000001, step=0.01)
+                }
+
+        # Enhance button
+        if st.button("Enhance Image"):
+            with st.spinner("Enhancing image..."):
+                enhanced_image = np.zeros_like(image_float)
+                for c in range(3):
+                    enhanced_image[:, :, c] = enhance_channel(
+                        image_float[:, :, c], depth,
+                        params[c]["veil"],
+                        params[c]["backscatter"],
+                        params[c]["recover"],
+                        params[c]["attenuation"]
+                    )
+                enhanced_image = gray_world_white_balance(enhanced_image)
+                enhanced_image = np.clip(enhanced_image, 0, 1)
+
+            # Convert enhanced image for display and download
+            enhanced_image_uint8 = (enhanced_image * 255).astype(np.uint8)
+            enhanced_pil = Image.fromarray(cv2.cvtColor(enhanced_image_uint8, cv2.COLOR_RGB2BGR))
+
+            # Display enhanced image
+            st.image(enhanced_pil, caption="Enhanced Image", use_column_width=True)
+
+            # Download button
+            buf = io.BytesIO()
+            enhanced_pil.save(buf, format="PNG")
+            byte_im = buf.getvalue()
+            st.download_button(
+                label="Download Enhanced Image",
+                data=byte_im,
+                file_name="enhanced_image.png",
+                mime="image/png"
+            )
+
+if __name__ == "__main__":
+    main()