Hashredacted/oil_spill
0
1# app.py2import streamlit as st3import numpy as np4import pandas as pd5from PIL import Image6import tensorflow as tf7from tensorflow.keras.models import load_model8import io9import warnings10 11# Suppress warnings12warnings.filterwarnings('ignore')13 14# Set page configuration15st.set_page_config(16 page_title="Oil Spill Segmentation",17 page_icon="🛢️",18 layout="wide"19)20 21# Constants22COLOR_MAP = [23 [0, 0, 0], [0, 255, 255], [255, 0, 0], [153, 76, 0], [0, 153, 0]24]25 26CLASS_NAMES = ["Background", "Oil Spill", "Ship", "Land", "Vegetation"]27 28@st.cache_resource29def load_model_cached():30 """Load the trained U-Net model"""31 try:32 return load_model("src/unet_model.keras")33 except:34 st.error("Could not load model. Ensure 'unet_model.keras' is in the directory.")35 return None36 37def preprocess_image(image):38 """Preprocess image for model input"""39 # Convert to RGB if needed40 if image.mode != 'RGB':41 image = image.convert('RGB')42 43 # Resize and normalize44 image_resized = np.array(image.resize((256, 256)))45 image_normalized = image_resized / 255.046 return np.expand_dims(image_normalized, axis=0), image_resized47 48def postprocess_mask(prediction):49 """Convert model prediction to mask"""50 mask = np.argmax(prediction, axis=-1)51 colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)52 53 for class_idx, color in enumerate(COLOR_MAP):54 colored_mask[mask == class_idx] = color55 56 return mask[0], colored_mask[0]57 58def create_simple_overlay(original, mask):59 """Simple overlay using PIL"""60 if isinstance(mask, np.ndarray):61 mask_pil = Image.fromarray(mask)62 else:63 mask_pil = mask64 65 # Resize original to match mask size (256x256)66 original_resized = original.resize((256, 256))67 68 # Create overlay by blending69 overlay = Image.blend(original_resized.convert('RGBA'), 70 mask_pil.convert('RGBA'), 71 alpha=0.5)72 73 return overlay.convert('RGB')74 75def pil_to_bytes(image):76 """Convert PIL image to bytes for download"""77 img_byte_arr = io.BytesIO()78 image.save(img_byte_arr, format='PNG')79 img_byte_arr.seek(0)80 return img_byte_arr81 82def main():83 st.title("🛢️ Oil Spill Segmentation")84 st.write("Upload a satellite image to detect oil spills and other features.")85 86 # Model loading87 model = load_model_cached()88 if model is None:89 return90 91 # File upload92 uploaded_file = st.file_uploader("Choose image", type=['jpg', 'jpeg', 'png', 'tif', 'tiff'])93 94 if uploaded_file is not None:95 try:96 # Load image97 image = Image.open(uploaded_file)98 99 col1, col2 = st.columns(2)100 101 with col1:102 st.subheader("Original Image")103 st.image(image, use_container_width=True)104 105 # Process image106 with st.spinner("Analyzing image..."):107 processed, image_resized = preprocess_image(image)108 prediction = model.predict(processed, verbose=0)109 class_mask, colored_mask = postprocess_mask(prediction)110 111 # Convert to PIL images for display and download112 mask_pil = Image.fromarray(colored_mask)113 overlay = create_simple_overlay(image, colored_mask)114 115 # Display results116 with col2:117 st.subheader("Segmentation Mask")118 st.image(mask_pil, use_container_width=True)119 120 st.subheader("Overlay")121 st.image(overlay, use_container_width=True)122 123 # Analysis124 st.subheader("Analysis")125 total_pixels = class_mask.size126 127 # Calculate all class percentages128 class_percentages = []129 for i in range(5):130 class_pixels = np.sum(class_mask == i)131 percentage = (class_pixels / total_pixels) * 100132 class_percentages.append(percentage)133 134 oil_percent = class_percentages[1]135 136 # Overall metrics137 col1, col2, col3 = st.columns(3)138 col1.metric("Total Pixels", f"{total_pixels:,}")139 col2.metric("Oil Spill Coverage", f"{oil_percent:.2f}%")140 141 with col3:142 if oil_percent > 5:143 st.error("🚨 Major oil spill detected!")144 elif oil_percent > 1:145 st.warning("⚠️ Significant oil spill detected")146 elif oil_percent > 0.1:147 st.info("ℹ️ Minor oil spill detected")148 else:149 st.success("✅ No significant oil spill detected")150 151 # Detailed class breakdown152 st.subheader("Detailed Class Analysis")153 154 # Create a clean table for class distribution155 class_data = []156 for idx, (class_name, color) in enumerate(zip(CLASS_NAMES, COLOR_MAP)):157 pixels = np.sum(class_mask == idx)158 percentage = class_percentages[idx]159 class_data.append({160 'Class': class_name,161 'Pixels': f"{pixels:,}",162 'Percentage': f"{percentage:.2f}%",163 'Color': f"rgb({color[0]}, {color[1]}, {color[2]})"164 })165 166 # Display as metrics in columns167 st.write("**Class Distribution:**")168 metric_cols = st.columns(5)169 for idx, data in enumerate(class_data):170 with metric_cols[idx]:171 # Create color swatch172 color_array = np.full((30, 30, 3), COLOR_MAP[idx], dtype=np.uint8)173 st.image(color_array, use_container_width=True)174 st.metric(175 label=data['Class'],176 value=data['Percentage'],177 help=f"{data['Pixels']} pixels"178 )179 180 # Optional: Display as table181 with st.expander("View Detailed Table"):182 df = pd.DataFrame(class_data)183 st.dataframe(df[['Class', 'Pixels', 'Percentage']], use_container_width=True)184 185 # Download options186 st.subheader("Download Results")187 col1, col2 = st.columns(2)188 189 with col1:190 mask_bytes = pil_to_bytes(mask_pil)191 st.download_button(192 "📥 Download Mask",193 data=mask_bytes,194 file_name="segmentation_mask.png",195 mime="image/png",196 use_container_width=True197 )198 199 with col2:200 overlay_bytes = pil_to_bytes(overlay)201 st.download_button(202 "📥 Download Overlay", 203 data=overlay_bytes,204 file_name="segmentation_overlay.png",205 mime="image/png",206 use_container_width=True207 )208 209 except Exception as e:210 st.error(f"Error processing image: {str(e)}")211 st.info("Try uploading a different image format (JPEG or PNG recommended)")212 213if __name__ == "__main__":214 main()