add codebase

app.py

@@ -0,0 +1,143 @@
+import streamlit as st
+import cv2
+import numpy as np
+import os
+from PIL import Image
+import torch
+from predict import load_onnx_model
+from utils.helpers import calculate_deforestation_metrics, create_overlay
+
+torch.classes.__path__ = []
+
+# Set page config
+st.set_page_config(page_title="Deforestation Detection", page_icon="🌳", layout="wide")
+
+# Set constants
+MODEL_INPUT_SIZE = 256  # The size our model expects
+
+# Load ONNX model
+@st.cache_resource
+def load_cached_onnx_model():
+    model_path = "models/deforestation_model.onnx"
+    return load_onnx_model(model_path, input_size=MODEL_INPUT_SIZE)
+
+def process_image(model, image):
+    """Process a single image and return results"""
+    # Save original image dimensions for display
+    orig_height, orig_width = image.shape[:2]
+
+    # Make prediction
+    mask = model.predict(image)
+
+    # Resize mask back to original dimensions for display
+    display_mask = cv2.resize(mask, (orig_width, orig_height))
+
+    # Create binary mask for visualization
+    binary_mask = (display_mask > 0.5).astype(np.uint8) * 255
+
+    # Create colored overlay
+    overlay = create_overlay(cv2.cvtColor(image, cv2.COLOR_BGR2RGB), display_mask)
+
+    # Calculate metrics
+    metrics = calculate_deforestation_metrics(mask)
+
+    return binary_mask, overlay, metrics
+
+def main():
+    # App title and description
+    st.title("🌳 Deforestation Detection")
+    st.markdown(
+        """
+    This app detects areas of deforestation in satellite or aerial images of forests.
+    Upload an image to get started!
+    """
+    )
+
+    # Model info
+    st.info(
+        f"⚙️ Model optimized for {MODEL_INPUT_SIZE}x{MODEL_INPUT_SIZE} pixel images using ONNX runtime"
+    )
+
+    # Load model
+    try:
+        model = load_cached_onnx_model()
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        st.info(
+            "Make sure you have converted your PyTorch model to ONNX format using the utils/onnx_converter.py script."
+        )
+        st.code(
+            "python -m utils.onnx_converter models/best_model_100.pth models/deforestation_model.onnx"
+        )
+        return
+
+    # File uploader for images
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+
+    if uploaded_file is not None:
+        # Load image
+        file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
+        image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
+
+        # Display original image
+        st.subheader("Original Image")
+        st.image(
+            cv2.cvtColor(image, cv2.COLOR_BGR2RGB),
+            caption="Uploaded Image",
+            use_container_width=True,
+        )
+
+        # Add a spinner while processing
+        with st.spinner("Processing..."):
+            try:
+                # Process image
+                binary_mask, overlay, metrics = process_image(model, image)
+
+                # Display results in columns
+                col1, col2 = st.columns(2)
+
+                with col1:
+                    st.subheader("Segmentation Result")
+                    st.image(
+                        binary_mask,
+                        caption="Forest Areas (White)",
+                        use_container_width=True,
+                    )
+
+                with col2:
+                    st.subheader("Overlay Visualization")
+                    st.image(
+                        overlay,
+                        caption="Green: Forest, Brown: Deforested",
+                        use_container_width=True,
+                    )
+
+                # Display metrics
+                st.subheader("Deforestation Analysis")
+
+                # Create metrics cards
+                metrics_col1, metrics_col2, metrics_col3 = st.columns(3)
+
+                with metrics_col1:
+                    st.metric(
+                        label="Forest Coverage",
+                        value=f"{metrics['forest_percentage']:.1f}%",
+                    )
+
+                with metrics_col2:
+                    st.metric(
+                        label="Deforested Area",
+                        value=f"{metrics['deforested_percentage']:.1f}%",
+                    )
+
+                with metrics_col3:
+                    st.metric(
+                        label="Deforestation Level",
+                        value=metrics["deforestation_level"],
+                    )
+
+            except Exception as e:
+                st.error(f"Error during processing: {e}")
+
+if __name__ == "__main__":
+    main()

predict.py

@@ -0,0 +1,165 @@
+import torch
+import numpy as np
+import cv2
+import onnxruntime as ort
+from utils.preprocess import preprocess_image
+from utils.model import load_model
+
+
+class PredictionEngine:
+    def __init__(self, model_path=None, use_onnx=True, input_size=256):
+        """
+        Initialize the prediction engine
+
+        Args:
+            model_path: Path to the model file (PyTorch or ONNX)
+            use_onnx: Whether to use ONNX runtime for inference
+            input_size: Input size for the model (default is 256)
+        """
+        self.use_onnx = use_onnx
+        self.input_size = input_size
+
+        if model_path:
+            if use_onnx:
+                self.model = self._load_onnx_model(model_path)
+            else:
+                self.model = load_model(model_path)
+        else:
+            self.model = None
+
+    def _load_onnx_model(self, model_path):
+        """
+        Load an ONNX model
+
+        Args:
+            model_path: Path to the ONNX model
+
+        Returns:
+            ONNX Runtime InferenceSession
+        """
+        # Try with CUDA first, fall back to CPU if needed
+        try:
+            session = ort.InferenceSession(
+                model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"]
+            )
+            print("ONNX model loaded with CUDA support")
+            return session
+        except Exception as e:
+            print(f"Could not load ONNX model with CUDA, falling back to CPU: {e}")
+            session = ort.InferenceSession(
+                model_path, providers=["CPUExecutionProvider"]
+            )
+            print("ONNX model loaded with CPU support")
+            return session
+
+    def preprocess(self, image):
+        """
+        Preprocess an image for prediction
+
+        Args:
+            image: Input image (numpy array)
+
+        Returns:
+            Processed image suitable for the model
+        """
+        # Keep the original image for reference
+        self.original_shape = image.shape[:2]
+
+        # Preprocess image
+        if self.use_onnx:
+            # For ONNX, we need to ensure the input is exactly the expected size
+            tensor = preprocess_image(image, img_size=self.input_size)
+            return tensor.numpy()
+        else:
+            # For PyTorch
+            return preprocess_image(image, img_size=self.input_size)
+
+    def predict(self, image):
+        """
+        Make a prediction on an image
+
+        Args:
+            image: Input image (numpy array)
+
+        Returns:
+            Predicted mask
+        """
+        if self.model is None:
+            raise ValueError("Model not loaded. Initialize with a valid model path.")
+
+        # Preprocess the image
+        processed_input = self.preprocess(image)
+
+        # Run inference
+        if self.use_onnx:
+            # Get input and output names
+            input_name = self.model.get_inputs()[0].name
+            output_name = self.model.get_outputs()[0].name
+
+            # Run ONNX inference
+            outputs = self.model.run([output_name], {input_name: processed_input})
+
+            # Apply sigmoid to output
+            mask = 1 / (1 + np.exp(-outputs[0].squeeze()))
+        else:
+            # PyTorch inference
+            with torch.no_grad():
+                # Move to device
+                device = next(self.model.parameters()).device
+                processed_input = processed_input.to(device)
+
+                # Forward pass
+                output = self.model(processed_input)
+                output = torch.sigmoid(output)
+
+                # Convert to numpy
+                mask = output.cpu().numpy().squeeze()
+
+        return mask
+
+
+def load_pytorch_model(model_path):
+    """
+    Load the PyTorch model for prediction
+
+    Args:
+        model_path: Path to the PyTorch model
+
+    Returns:
+        PredictionEngine instance
+    """
+    return PredictionEngine(model_path, use_onnx=False)
+
+
+def load_onnx_model(model_path, input_size=256):
+    """
+    Load the ONNX model for prediction
+
+    Args:
+        model_path: Path to the ONNX model
+        input_size: Input size for the model
+
+    Returns:
+        PredictionEngine instance
+    """
+    return PredictionEngine(model_path, use_onnx=True, input_size=input_size)
+
+
+# For backwards compatibility
+def predict(model, image):
+    """
+    Legacy function for prediction
+
+    Args:
+        model: Model instance
+        image: Input image
+
+    Returns:
+        Predicted mask
+    """
+    if isinstance(model, PredictionEngine):
+        return model.predict(image)
+
+    engine = PredictionEngine(use_onnx=True)
+    engine.model = model
+    return engine.predict(image)

requirements.txt

@@ -0,0 +1,13 @@
+streamlit
+torch
+torchvision
+opencv-python
+albumentations
+numpy
+Pillow
+scikit-image
+scikit-learn
+matplotlib
+onnxruntime
+onnxruntime-gpu
+onnx

utils/helpers.py

@@ -0,0 +1,73 @@
+import cv2
+import numpy as np
+
+
+def calculate_deforestation_metrics(mask, threshold=0.5):
+    """
+    Calculate deforestation metrics from the predicted mask
+
+    Args:
+        mask: Predicted mask (numpy array)
+        threshold: Threshold to binarize the mask (default is 0.5)
+
+    Returns:
+        Dictionary containing deforestation metrics
+    """
+    # Binarize the mask
+    binary_mask = (mask > threshold).astype(np.uint8)
+
+    # Calculate pixel counts
+    total_pixels = binary_mask.size
+    forest_pixels = np.sum(binary_mask)
+    deforested_pixels = total_pixels - forest_pixels
+
+    # Calculate percentages
+    forest_percentage = (forest_pixels / total_pixels) * 100
+    deforested_percentage = (deforested_pixels / total_pixels) * 100
+
+    # Determine deforestation level
+    if deforested_percentage < 20:
+        level = "Low"
+    elif deforested_percentage < 50:
+        level = "Medium"
+    else:
+        level = "High"
+
+    return {
+        "forest_pixels": forest_pixels,
+        "deforested_pixels": deforested_pixels,
+        "forest_percentage": forest_percentage,
+        "deforested_percentage": deforested_percentage,
+        "deforestation_level": level,
+    }
+
+
+def create_overlay(original_image, mask, threshold=0.5, alpha=0.5):
+    """
+    Create a visualization by overlaying the mask on the original image
+
+    Args:
+        original_image: Original RGB image
+        mask: Predicted mask
+        threshold: Threshold to binarize the mask
+        alpha: Opacity of the overlay
+
+    Returns:
+        Overlay image
+    """
+    # Resize mask to match original image if needed
+    if original_image.shape[:2] != mask.shape[:2]:
+        mask = cv2.resize(mask, (original_image.shape[1], original_image.shape[0]))
+
+    # Create binary mask
+    binary_mask = (mask > threshold).astype(np.uint8) * 255
+
+    # Create a colored mask (green for forest, red for deforested)
+    colored_mask = np.zeros_like(original_image)
+    colored_mask[binary_mask == 255] = [0, 255, 0]  # Green for forest
+    colored_mask[binary_mask == 0] = [150, 75, 0]  # Brown for deforested
+
+    # Create overlay
+    overlay = cv2.addWeighted(original_image, 1 - alpha, colored_mask, alpha, 0)
+
+    return overlay

utils/model.py

@@ -0,0 +1,143 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+
+class AttentionGate(nn.Module):
+    def __init__(self, F_g, F_l, F_int):
+        super(AttentionGate, self).__init__()
+        self.W_g = nn.Sequential(
+            nn.Conv2d(F_g, F_int, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(F_int),
+        )
+
+        self.W_x = nn.Sequential(
+            nn.Conv2d(F_l, F_int, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(F_int),
+        )
+
+        self.psi = nn.Sequential(
+            nn.Conv2d(F_int, 1, kernel_size=1, stride=1, padding=0, bias=True),
+            nn.BatchNorm2d(1),
+            nn.Sigmoid(),
+        )
+
+        self.relu = nn.ReLU(inplace=True)
+
+    def forward(self, g, x):
+        g1 = self.W_g(g)
+        x1 = self.W_x(x)
+        psi = self.relu(g1 + x1)
+        psi = self.psi(psi)
+
+        return x * psi
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super(DoubleConv, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class AttentionUNet(nn.Module):
+    def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):
+        super(AttentionUNet, self).__init__()
+        self.ups = nn.ModuleList()
+        self.downs = nn.ModuleList()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.attention_gates = nn.ModuleList()
+
+        # Down part of U-Net
+        for feature in features:
+            self.downs.append(DoubleConv(in_channels, feature))
+            in_channels = feature
+
+        # Up part of U-Net
+        for feature in reversed(features):
+            self.ups.append(
+                nn.ConvTranspose2d(
+                    feature * 2,
+                    feature,
+                    kernel_size=2,
+                    stride=2,
+                )
+            )
+            # Attention Gate
+            self.attention_gates.append(
+                AttentionGate(F_g=feature, F_l=feature, F_int=feature // 2)
+            )
+
+            self.ups.append(DoubleConv(feature * 2, feature))
+
+        # Bottleneck
+        self.bottleneck = DoubleConv(features[-1], features[-1] * 2)
+
+        # Final Conv
+        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)
+
+    def forward(self, x):
+        skip_connections = []
+
+        # Encoder path
+        for down in self.downs:
+            x = down(x)
+            skip_connections.append(x)
+            x = self.pool(x)
+
+        x = self.bottleneck(x)
+        skip_connections = skip_connections[::-1]  # Reverse to use from back
+
+        # Decoder path
+        for idx in range(0, len(self.ups), 2):
+            x = self.ups[idx](x)
+            skip_connection = skip_connections[idx // 2]
+
+            # If sizes don't match
+            if x.shape != skip_connection.shape:
+                x = F.interpolate(x, size=skip_connection.shape[2:])
+
+            # Apply attention gate
+            skip_connection = self.attention_gates[idx // 2](g=x, x=skip_connection)
+
+            # Concatenate
+            concat_skip = torch.cat((skip_connection, x), dim=1)
+            x = self.ups[idx + 1](concat_skip)
+
+        # Final conv
+        return self.final_conv(x)
+
+
+def load_model(model_path):
+    """
+    Load the trained model
+
+    Args:
+        model_path: Path to the model weights
+
+    Returns:
+        Loaded model
+    """
+    # Define device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Initialize model
+    model = AttentionUNet(in_channels=3, out_channels=1)
+
+    # Load model weights
+    model.load_state_dict(torch.load(model_path, map_location=device))
+
+    # Set model to evaluation mode
+    model.eval()
+
+    return model.to(device)

utils/onnx_converter.py

@@ -0,0 +1,99 @@
+import torch
+import sys
+import numpy as np
+from utils.model import AttentionUNet
+import onnx
+import onnxruntime as ort
+
+
+def convert_to_onnx(pytorch_model_path, onnx_output_path, input_size=256):
+    """
+    Convert a PyTorch model to ONNX format
+
+    Args:
+        pytorch_model_path: Path to the PyTorch model
+        onnx_output_path: Path to save the ONNX model
+        input_size: Input size for the model (default is 256x256)
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device used for conversion: {device}")
+
+    model = AttentionUNet(in_channels=3, out_channels=1)
+    model.to(device)
+    model.load_state_dict(torch.load(pytorch_model_path, map_location=device))
+
+    model.eval()
+
+    # Create dummy input
+    dummy_input = torch.randn(1, 3, input_size, input_size, device=device)
+
+    # Export the model
+    torch.onnx.export(
+        model,  # model being run
+        dummy_input,  # model input (or a tuple for multiple inputs)
+        onnx_output_path,  # where to save the model
+        export_params=True,  # store the trained parameter weights inside the model file
+        opset_version=12,  # the ONNX version to export the model to
+        do_constant_folding=True,  # whether to execute constant folding for optimization
+        input_names=["input"],
+        output_names=["output"],
+        dynamic_axes={
+            "input": {0: "batch_size", 2: "height", 3: "width"},  # variable length axes
+            "output": {0: "batch_size", 2: "height", 3: "width"},
+        },
+    )
+
+    print(f"Model converted and saved to {onnx_output_path}")
+
+    verify_onnx_model(onnx_output_path, input_size)
+
+
+def verify_onnx_model(onnx_model_path, input_size=256):
+    """
+    Verify the ONNX model to ensure it was exported correctly
+
+    Args:
+        onnx_model_path: Path to the ONNX model
+        input_size: Input size used during export
+    """
+    try:
+        onnx_model = onnx.load(onnx_model_path)
+        onnx.checker.check_model(onnx_model)
+        print("ONNX model is valid")
+    except Exception as e:
+        print(f"ONNX model validation failed: {e}")
+        return False
+
+    try:
+        session = ort.InferenceSession(
+            onnx_model_path, providers=["CPUExecutionProvider"]
+        )
+
+        input_data = np.random.rand(1, 3, input_size, input_size).astype(np.float32)
+
+        # Get input and output names
+        input_name = session.get_inputs()[0].name
+        output_name = session.get_outputs()[0].name
+
+        # Run inference
+        outputs = session.run([output_name], {input_name: input_data})
+
+        print(f"ONNX model inference test passed. Output shape: {outputs[0].shape}")
+        return True
+    except Exception as e:
+        print(f"ONNX model inference test failed: {e}")
+        return False
+
+
+if __name__ == "__main__":
+    if len(sys.argv) < 3:
+        print(
+            "Usage: python -m utils.onnx_converter <pytorch_model_path> <onnx_output_path> [input_size]"
+        )
+        sys.exit(1)
+
+    pytorch_model_path = sys.argv[1]
+    onnx_output_path = sys.argv[2]
+    input_size = int(sys.argv[3]) if len(sys.argv) > 3 else 256
+
+    convert_to_onnx(pytorch_model_path, onnx_output_path, input_size)

utils/preprocess.py

@@ -0,0 +1,36 @@
+import cv2
+import torch
+import numpy as np
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+def preprocess_image(image, img_size=256):
+    """
+    Preprocess the input image for model prediction
+
+    Args:
+        image: Input image (numpy array)
+        img_size: Size to resize image to (model was trained on 256x256)
+
+    Returns:
+        Preprocessed image tensor
+    """
+    # Define the transformation
+    transform = A.Compose(
+        [
+            A.Resize(height=img_size, width=img_size),
+            A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+            ToTensorV2(),
+        ]
+    )
+
+    # Apply the transformation
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    augmented = transform(image=image)
+    image_tensor = augmented["image"]
+
+    # Add batch dimension
+    image_tensor = image_tensor.unsqueeze(0)
+
+    return image_tensor