Initial commit

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

Readme.md

@@ -0,0 +1,129 @@
+---
+language: "en"
+license: "gpl-3.0"
+tags:
+  - segmentation
+  - computer-vision
+  - yolo
+  - beach
+  - water
+  - open-source
+task_categories:
+  - image-segmentation
+---
+
+# 🌊 Water Surface Segmentation on Beach Images
+
+## Model Overview
+This model performs **semantic segmentation of water surfaces** in beach or coastal images.
+It’s a fine-tuned version of **YOLOv11n**, adapted for **binary segmentation** with a single class: **`water`**.
+
+Built for lightweight, real-time deployment, the model achieves strong accuracy while remaining small and efficient.
+
+---
+
+## 🧠 Model Details
+- **Architecture**: YOLOv11n segmentation head (binary)
+- **Base framework**: PyTorch / Ultralytics YOLOv11
+- **Input size**: 640×640 RGB images  
+- **Output**: Binary segmentation mask (1 class — `water`)
+- **Model file**: `nwsd-v2.pt` (≈6 MB)
+
+---
+
+## 🚀 Key Features
+- ⚡ **Real-time inference** on CPU/GPU  
+- 🖼 **Outputs**: Binary masks, overlays, and coverage statistics  
+- 📊 **Evaluation tools** included for metrics & visualization  
+- 🐍 **Easy Python integration** via a simple API (`nwsd_api.py`)
+
+---
+
+## 📈 Performance
+| Metric | Value | Notes |
+|:--|:--|:--|
+| **mAP50** | > 0.85 | On validation set |
+| **Inference speed** | ~50 ms/image | On CPU |
+| **GPU memory** | < 2 GB | For 640×640 input |
+
+---
+
+## 🗂 Dataset
+- **Type**: Binary segmentation  
+- **Classes**: `water`  
+- **Annotations**: PNG masks  
+- **Source**: Custom-labeled beach dataset  
+
+🔗 [Dataset on Roboflow](https://universe.roboflow.com/neptune-uxxqf/neptune-water-surface-detection)
+
+---
+
+## 🧩 Intended Uses
+**Use cases:**
+- Coastal or maritime monitoring  
+- Beach safety & drowning prevention systems  
+- Environmental analysis (e.g., water coverage estimation)  
+
+**Limitations:**
+- Designed for daylight, clear beach imagery  
+- May underperform in low-visibility or night-time scenes
+
+---
+
+## 🧪 How to Use
+
+### Load model from Hub
+```python
+from huggingface_hub import hf_hub_download
+import torch
+
+model_path = hf_hub_download(repo_id="Ehlum-Lucas/NWSD", filename="nwsd-v2.pt")
+model = torch.load(model_path, map_location="cpu")
+model.eval()
+```
+
+### Inference example
+```python
+from PIL import Image
+import torch
+from torchvision import transforms
+
+img = Image.open("beachTest.jpg").convert("RGB")
+input_tensor = transforms.ToTensor()(img).unsqueeze(0)
+
+with torch.no_grad():
+    pred = model(input_tensor)
+```
+
+### ⚙️ Training
+You can fine-tune or retrain the model using YOLOv11 tools:
+```bash
+python train.py --data data.yaml --weights <path_to_weights> --img 640 --batch 16 --epochs 50
+```
+Example configuration (data.yaml) defines paths to your datasets and class names.
+
+### 🧭 Evaluation
+```bash
+python evaluate.py --data data.yaml --weights model/nwsd-v2.pt
+```
+
+Generates:
+
+- Binary mask
+- Overlay visualization
+- Water coverage stats
+
+## License
+This model is released under the **GPL-3.0 License**. See the [LICENSE](LICENSE) file for details.
+
+## Citation
+If you use this model in your work, please consider citing:
+```latex
+@misc{nwsd2025,
+  title={Water Surface Segmentation on Beach Images},
+  author={Lucas Iglesia},
+  year={2025},
+  howpublished={\url{https://huggingface.co/Ehlum-Lucas/NWSD}}
+}
+```
+

decode_mask.py

@@ -0,0 +1,4 @@
+import base64
+
+with open("mask.png", "wb") as f:
+    f.write(base64.b64decode("iVBORw0KGgoAAAANSUhEUgAAAoAAAAHgCAAAAAAQuoM4AAAFHElEQVR4Ae3BUbITBBQFwZn9L/pafqCUUvrgBU5CpluSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGJBmSZEiSIUmGPF6Z5LV5vDLJa/N4ZZLX5vHKJK/N4zcmeXIeb0X+dvyT/K/ja5LP8ciDSL6bRx5O8kEe+Zkk/8UjC5I/eeS5yDvxyKuT1+WR35U8P4+8DXk6Hnl38sXxEfI4HsnnyY/xSH4x+YtHsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY5HsuOR7HgkOx7Jjkey45HseCQ7HsmOR7Ljkex4JDseyY7AkYzIhx3Jo8knHMnnyOMcyXeSX+NIvkGmjrw3eSbHA8k3HXkikg858jNIfsSRh5A82pEPk+wdb0vybI43Inkxx+9E8sKOr8g3Hc9M8j6OpyN5Q8ezkLyp4xlI3tixJnl7x44kXxy/nCT/dvwiknzE8VNI8j2Oh5LkBxyPIcknHJ8jyQMcP0aSBzq+jyRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgxJMiTJkCRDkgz9AV/Pnlo/ledrAAAAAElFTkSuQmCC"))

evaluate.py

@@ -0,0 +1,165 @@
+#!/usr/bin/env python3
+"""
+Water Surface Segmentation Evaluation Script
+Evaluate the trained model on a validation dataset.
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from ultralytics import YOLO
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Evaluate water surface segmentation model",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--data",
+        type=str,
+        required=True,
+        help="Path to validation dataset or data.yaml file"
+    )
+
+    parser.add_argument(
+        "--weights",
+        type=str,
+        default="model/nwsd-v2.pt",
+        help="Path to model weights file"
+    )
+
+    parser.add_argument(
+        "--img",
+        type=int,
+        default=640,
+        help="Image size for evaluation"
+    )
+
+    parser.add_argument(
+        "--batch",
+        type=int,
+        default=16,
+        help="Batch size for evaluation"
+    )
+
+    parser.add_argument(
+        "--conf",
+        type=float,
+        default=0.25,
+        help="Confidence threshold"
+    )
+
+    parser.add_argument(
+        "--iou",
+        type=float,
+        default=0.45,
+        help="IoU threshold for NMS"
+    )
+
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="",
+        help="Device to use for evaluation (cpu, cuda, mps)"
+    )
+
+    parser.add_argument(
+        "--project",
+        type=str,
+        default="runs/segment",
+        help="Project directory for results"
+    )
+
+    parser.add_argument(
+        "--name",
+        type=str,
+        default="nwsd_eval",
+        help="Experiment name"
+    )
+
+    parser.add_argument(
+        "--save-json",
+        action="store_true",
+        help="Save results in JSON format"
+    )
+
+    parser.add_argument(
+        "--save-txt",
+        action="store_true",
+        help="Save results in TXT format"
+    )
+
+    parser.add_argument(
+        "--plots",
+        action="store_true",
+        help="Generate evaluation plots"
+    )
+
+    return parser.parse_args()
+
+
+def validate_inputs(args: argparse.Namespace) -> None:
+    """Validate input arguments."""
+    if not os.path.exists(args.data):
+        raise FileNotFoundError(f"Data path not found: {args.data}")
+
+    if not os.path.exists(args.weights):
+        raise FileNotFoundError(f"Model weights not found: {args.weights}")
+
+
+def main():
+    """Main evaluation function."""
+    args = parse_arguments()
+
+    try:
+        validate_inputs(args)
+
+        print(f"Loading model: {args.weights}")
+        model = YOLO(args.weights)
+
+        eval_params = {
+            'data': args.data,
+            'imgsz': args.img,
+            'batch': args.batch,
+            'conf': args.conf,
+            'iou': args.iou,
+            'device': args.device,
+            'project': args.project,
+            'name': args.name,
+            'save_json': args.save_json,
+            'save_txt': args.save_txt,
+            'plots': args.plots,
+            'verbose': True,
+        }
+
+        print("Starting evaluation with parameters:")
+        for key, value in eval_params.items():
+            print(f"  {key}: {value}")
+
+        results = model.val(**eval_params)
+
+        print("\n" + "="*50)
+        print("EVALUATION RESULTS SUMMARY")
+        print("="*50)
+
+        if hasattr(results, 'box') and results.box is not None:
+            print(f"mAP50: {results.box.map50:.4f}")
+            print(f"mAP50-95: {results.box.map:.4f}")
+
+        if hasattr(results, 'seg') and results.seg is not None:
+            print(f"Segmentation mAP50: {results.seg.map50:.4f}")
+            print(f"Segmentation mAP50-95: {results.seg.map:.4f}")
+
+        print("\nEvaluation completed successfully!")
+
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

model_card_template.yaml

@@ -0,0 +1,87 @@
+# model_card_template.yaml
+
+# =====================================================
+# 🌊 Water Surface Segmentation on Beach Images
+# =====================================================
+# Hugging Face model metadata file
+# =====================================================
+
+language:
+  - en
+
+license: gpl-3.0
+
+library_name: pytorch
+
+tags:
+  - segmentation
+  - computer-vision
+  - yolo
+  - beach
+  - water
+  - open-source
+
+task_categories:
+  - image-segmentation
+
+model-index:
+  - name: Water Surface Segmentation (NWSD)
+    results:
+      - task:
+          type: image-segmentation
+          name: Image Segmentation
+        metrics:
+          - type: mAP50
+            name: Mean Average Precision @ 0.5
+            value: 0.85
+          - type: inference_speed
+            name: Inference Speed (CPU)
+            value: 50
+            unit: "ms/image"
+
+model_details:
+  description: >
+    A YOLOv11n-based segmentation model fine-tuned for detecting and segmenting
+    water surfaces in coastal or beach images. Trained on a custom-labeled dataset
+    containing a single class: "water".
+  developed_by: Lucas Iglesia
+  repo: https://huggingface.co/Lucas-Iglesia/NWSD
+  license: GPL-3.0
+  framework: PyTorch
+  model_size: 6.07 MB
+  input_size: "640x640"
+  num_classes: 1
+  class_labels: ["water"]
+  release_date: "2025-11-07"
+
+inference:
+  parameters:
+    device: "cpu or cuda"
+    conf_threshold: 0.5
+  example_inputs:
+    - beachTest.jpg
+  example_outputs:
+    - binary_mask.png
+    - overlay.png
+  usage_snippet: |
+    from huggingface_hub import hf_hub_download
+    import torch
+    model_path = hf_hub_download(repo_id="Ehlum-Lucas/NWSD", filename="nwsd-v2.pt")
+    model = torch.load(model_path, map_location="cpu")
+    model.eval()
+
+recommended_use:
+  - Coastal monitoring
+  - Beach safety and drowning prevention
+  - Environmental water coverage analysis
+
+limitations:
+  - Optimized for daylight beach scenes
+  - May underperform in low-visibility or night images
+
+citation:
+  - type: misc
+    title: "Water Surface Segmentation on Beach Images"
+    author: "Lucas Iglesia"
+    year: 2025
+    url: "https://huggingface.co/Ehlum-Lucas/NWSD"

nwsd-v2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:643363b33e702713dc69f38146bdeb6f6a47b1c7c32d7593a58b0a5c7f9b4722
+size 6360093

nwsd_api.py

@@ -0,0 +1,233 @@
+#!/usr/bin/env python3
+"""
+NWSD API - Simple Python API for water surface detection
+This module provides a simple interface for water surface segmentation.
+"""
+
+import os
+import cv2
+import numpy as np
+from typing import Optional, Tuple, Dict, Union
+from pathlib import Path
+from ultralytics import YOLO
+
+
+class WaterSurfaceDetector:
+    """Water Surface Detection API using YOLOv11n."""
+
+    def __init__(self, weights_path: str = "model/nwsd-v2.pt", device: str = "cpu"):
+        """
+        Initialize the water surface detector.
+
+        Args:
+            weights_path: Path to model weights
+            device: Device to use for inference (cpu, cuda, mps)
+        """
+        self.weights_path = weights_path
+        self.device = device
+        self.model = None
+        self._load_model()
+
+    def _load_model(self):
+        """Load the YOLO model."""
+        if not os.path.exists(self.weights_path):
+            raise FileNotFoundError(f"Model weights not found: {self.weights_path}")
+
+        self.model = YOLO(self.weights_path)
+        self.model.to(self.device)
+
+    def detect(self,
+               image: Union[str, np.ndarray],
+               conf: float = 0.25,
+               iou: float = 0.45) -> Dict:
+        """
+        Detect water surfaces in an image.
+
+        Args:
+            image: Path to image file or numpy array
+            conf: Confidence threshold
+            iou: IoU threshold for NMS
+
+        Returns:
+            Dictionary containing detection results
+        """
+        if isinstance(image, str):
+            img_array = cv2.imread(image)
+            if img_array is None:
+                raise ValueError(f"Could not load image: {image}")
+            image_path = image
+        else:
+            img_array = image
+            image_path = None
+
+        results = self.model(image_path if image_path else img_array,
+                           conf=conf, iou=iou, verbose=False)
+
+        return self._process_results(results, img_array)
+
+    def _process_results(self, results, original_image: np.ndarray) -> Dict:
+        """Process YOLO results into structured output."""
+        h, w = original_image.shape[:2]
+
+        output = {
+            "detected": False,
+            "binary_mask": None,
+            "overlay": None,
+            "water_percentage": 0.0,
+            "water_pixels": 0,
+            "total_pixels": h * w,
+            "bounding_boxes": [],
+            "confidence_scores": []
+        }
+
+        if len(results) == 0 or results[0].masks is None:
+            return output
+
+        result = results[0]
+
+        masks = result.masks.data.cpu().numpy()
+
+        if len(masks) == 0:
+            return output
+
+        combined_mask = np.zeros((h, w), dtype=np.uint8)
+
+        for mask in masks:
+            resized_mask = cv2.resize(mask, (w, h))
+            combined_mask = np.maximum(combined_mask, (resized_mask > 0.5).astype(np.uint8))
+
+        binary_mask = combined_mask * 255
+
+        overlay = original_image.copy()
+        colored_mask = np.zeros_like(original_image)
+        colored_mask[binary_mask > 0] = [0, 0, 255]
+        overlay = cv2.addWeighted(overlay, 0.7, colored_mask, 0.3, 0)
+
+        water_pixels = np.sum(binary_mask > 0)
+        water_percentage = (water_pixels / (h * w)) * 100
+
+        if result.boxes is not None:
+            boxes = result.boxes.xyxy.cpu().numpy()
+            scores = result.boxes.conf.cpu().numpy()
+
+            output["bounding_boxes"] = boxes.tolist()
+            output["confidence_scores"] = scores.tolist()
+
+        output.update({
+            "detected": True,
+            "binary_mask": binary_mask,
+            "overlay": overlay,
+            "water_percentage": water_percentage,
+            "water_pixels": int(water_pixels)
+        })
+
+        return output
+
+    def detect_batch(self,
+                     image_paths: list,
+                     conf: float = 0.25,
+                     iou: float = 0.45) -> Dict:
+        """
+        Detect water surfaces in multiple images.
+
+        Args:
+            image_paths: List of paths to image files
+            conf: Confidence threshold
+            iou: IoU threshold for NMS
+
+        Returns:
+            Dictionary with results for each image
+        """
+        results = {}
+
+        for image_path in image_paths:
+            try:
+                result = self.detect(image_path, conf, iou)
+                results[image_path] = result
+            except Exception as e:
+                results[image_path] = {"error": str(e)}
+
+        return results
+
+    def save_results(self,
+                     results: Dict,
+                     output_dir: str,
+                     base_name: str,
+                     save_mask: bool = True,
+                     save_overlay: bool = True) -> Dict[str, str]:
+        """
+        Save detection results to files.
+
+        Args:
+            results: Results from detect() method
+            output_dir: Directory to save results
+            base_name: Base name for output files
+            save_mask: Whether to save binary mask
+            save_overlay: Whether to save overlay
+
+        Returns:
+            Dictionary with saved file paths
+        """
+        os.makedirs(output_dir, exist_ok=True)
+        saved_files = {}
+
+        if save_mask and results["binary_mask"] is not None:
+            mask_path = os.path.join(output_dir, f"{base_name}_mask.png")
+            cv2.imwrite(mask_path, results["binary_mask"])
+            saved_files["mask"] = mask_path
+
+        if save_overlay and results["overlay"] is not None:
+            overlay_path = os.path.join(output_dir, f"{base_name}_overlay.png")
+            cv2.imwrite(overlay_path, results["overlay"])
+            saved_files["overlay"] = overlay_path
+
+        return saved_files
+
+    def get_water_classification(self, percentage: float) -> str:
+        """Classify water coverage level."""
+        if percentage < 10:
+            return "minimal"
+        elif percentage < 30:
+            return "low"
+        elif percentage < 50:
+            return "moderate"
+        elif percentage < 70:
+            return "high"
+        else:
+            return "very_high"
+
+
+# Example usage
+def main():
+    """Example usage of the WaterSurfaceDetector API."""
+    print("🌊 NWSD API Example")
+    print("=" * 30)
+
+    detector = WaterSurfaceDetector()
+
+    # Look for test images
+    test_images = list(Path("..").glob("*.jpg"))
+
+    if not test_images:
+        print("No test images found")
+        return
+
+    test_image = str(test_images[0])
+    print(f"Processing: {test_image}")
+
+    results = detector.detect(test_image)
+
+    print(f"Water detected: {results['detected']}")
+    print(f"Water coverage: {results['water_percentage']:.2f}%")
+    print(f"Classification: {detector.get_water_classification(results['water_percentage'])}")
+
+    # Save results
+    if results['detected']:
+        output_dir = "api_results"
+        base_name = Path(test_image).stem
+        saved_files = detector.save_results(results, output_dir, base_name)
+        print(f"Results saved to: {saved_files}")
+
+
+if __name__ == "__main__":
+    main()

predict.py

@@ -0,0 +1,298 @@
+#!/usr/bin/env python3
+"""
+Water Surface Segmentation Inference Script
+This script performs inference on beach images to segment water surfaces using YOLOv11n.
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+import cv2
+import numpy as np
+from ultralytics import YOLO
+import matplotlib
+matplotlib.use('Agg')  # Use non-interactive backend
+import matplotlib.pyplot as plt
+from typing import Optional, Tuple, List
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Perform water surface segmentation on beach images",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--image",
+        type=str,
+        required=True,
+        help="Path to input image file"
+    )
+
+    parser.add_argument(
+        "--weights",
+        type=str,
+        default="model/nwsd-v2.pt",
+        help="Path to model weights file"
+    )
+
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="Output directory for results (default: same as input image)"
+    )
+
+    parser.add_argument(
+        "--conf",
+        type=float,
+        default=0.25,
+        help="Confidence threshold for segmentation"
+    )
+
+    parser.add_argument(
+        "--iou",
+        type=float,
+        default=0.45,
+        help="IoU threshold for NMS"
+    )
+
+    parser.add_argument(
+        "--save-overlay",
+        action="store_true",
+        help="Save overlay visualization"
+    )
+
+    parser.add_argument(
+        "--save-mask",
+        action="store_true",
+        help="Save binary mask"
+    )
+
+    parser.add_argument(
+        "--save-results",
+        action="store_true",
+        help="Save results visualization plot"
+    )
+
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cpu",
+        help="Device to use for inference (cpu, cuda, mps)"
+    )
+
+    return parser.parse_args()
+
+
+def validate_inputs(args: argparse.Namespace) -> None:
+    """Validate input arguments."""
+    if not os.path.exists(args.image):
+        raise FileNotFoundError(f"Input image not found: {args.image}")
+
+    if not os.path.exists(args.weights):
+        raise FileNotFoundError(f"Model weights not found: {args.weights}")
+
+    valid_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
+    image_ext = Path(args.image).suffix.lower()
+    if image_ext not in valid_extensions:
+        raise ValueError(f"Unsupported image format: {image_ext}")
+
+
+def load_model(weights_path: str, device: str = "cpu") -> YOLO:
+    """Load YOLO model."""
+    try:
+        model = YOLO(weights_path)
+        model.to(device)
+        print(f"Model loaded successfully from: {weights_path}")
+        print(f"Using device: {device}")
+        return model
+    except Exception as e:
+        raise RuntimeError(f"Failed to load model: {str(e)}")
+
+
+def preprocess_image(image_path: str) -> Tuple[np.ndarray, Tuple[int, int]]:
+    """Load and preprocess image."""
+    image = cv2.imread(image_path)
+    if image is None:
+        raise ValueError(f"Could not read image: {image_path}")
+
+    original_shape = image.shape[:2]  # (height, width)
+    return image, original_shape
+
+
+def postprocess_results(results, original_shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
+    """Extract masks and create binary mask."""
+    if len(results) == 0 or results[0].masks is None:
+        print("No water surface detected in the image")
+        return None, None
+
+    result = results[0]
+
+    masks = result.masks.data.cpu().numpy()  # Shape: (N, H, W)
+
+    binary_mask = np.zeros(original_shape, dtype=np.uint8)
+
+    if len(masks) > 0:
+        resized_masks = []
+        for mask in masks:
+            resized_mask = cv2.resize(mask, (original_shape[1], original_shape[0]))
+            resized_masks.append(resized_mask)
+
+        combined_mask = np.max(resized_masks, axis=0)
+        binary_mask = (combined_mask > 0.5).astype(np.uint8) * 255
+
+    return binary_mask, masks
+
+
+def create_overlay(image: np.ndarray, binary_mask: np.ndarray, alpha: float = 0.3) -> np.ndarray:
+    """Create overlay visualization."""
+    overlay = image.copy()
+
+    colored_mask = np.zeros_like(image)
+    colored_mask[binary_mask > 0] = [255, 0, 0]
+
+    overlay = cv2.addWeighted(overlay, 1 - alpha, colored_mask, alpha, 0)
+
+    return overlay
+
+
+def save_results(
+    image: np.ndarray,
+    binary_mask: Optional[np.ndarray],
+    overlay: Optional[np.ndarray],
+    output_dir: str,
+    base_name: str,
+    save_mask: bool = False,
+    save_overlay: bool = False
+) -> None:
+    """Save results to output directory."""
+    os.makedirs(output_dir, exist_ok=True)
+
+    if save_mask and binary_mask is not None:
+        mask_path = os.path.join(output_dir, f"{base_name}_mask.png")
+        cv2.imwrite(mask_path, binary_mask)
+        print(f"Binary mask saved to: {mask_path}")
+
+    if save_overlay and overlay is not None:
+        overlay_path = os.path.join(output_dir, f"{base_name}_overlay.png")
+        cv2.imwrite(overlay_path, overlay)
+        print(f"Overlay visualization saved to: {overlay_path}")
+
+
+def display_results(
+    image: np.ndarray,
+    binary_mask: Optional[np.ndarray],
+    overlay: Optional[np.ndarray],
+    output_dir: str = ".",
+    base_name: str = "result"
+) -> None:
+    """Display results using matplotlib."""
+    num_plots = 1 + (binary_mask is not None) + (overlay is not None)
+
+    plt.figure(figsize=(5 * num_plots, 5))
+
+    plot_idx = 1
+
+    plt.subplot(1, num_plots, plot_idx)
+    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
+    plt.title("Original Image")
+    plt.axis('off')
+    plot_idx += 1
+
+    if binary_mask is not None:
+        plt.subplot(1, num_plots, plot_idx)
+        plt.imshow(binary_mask, cmap='gray')
+        plt.title("Water Surface Mask")
+        plt.axis('off')
+        plot_idx += 1
+
+    if overlay is not None:
+        plt.subplot(1, num_plots, plot_idx)
+        plt.imshow(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
+        plt.title("Overlay Visualization")
+        plt.axis('off')
+
+    plt.tight_layout()
+
+    plot_path = os.path.join(output_dir, f"{base_name}_results.png")
+    plt.savefig(plot_path, dpi=150, bbox_inches='tight')
+    print(f"Results visualization saved to: {plot_path}")
+    plt.close()
+
+
+def calculate_water_percentage(binary_mask: np.ndarray) -> float:
+    """Calculate percentage of water surface in the image."""
+    if binary_mask is None:
+        return 0.0
+
+    total_pixels = binary_mask.shape[0] * binary_mask.shape[1]
+    water_pixels = np.sum(binary_mask > 0)
+
+    return (water_pixels / total_pixels) * 100
+
+
+def main():
+    """Main inference function."""
+    args = parse_arguments()
+
+    try:
+        validate_inputs(args)
+
+        if args.output is None:
+            output_dir = os.path.dirname(args.image)
+
+            if not output_dir:
+                output_dir = "."
+        else:
+            output_dir = args.output
+
+        base_name = Path(args.image).stem
+
+        model = load_model(args.weights, args.device)
+
+        image, original_shape = preprocess_image(args.image)
+
+        print(f"Processing image: {args.image}")
+        print(f"Image shape: {image.shape}")
+
+        results = model(
+            args.image,
+            conf=args.conf,
+            iou=args.iou,
+            verbose=False
+        )
+
+        binary_mask, masks = postprocess_results(results, original_shape)
+
+        overlay = None
+        if binary_mask is not None:
+            overlay = create_overlay(image, binary_mask)
+
+            water_percentage = calculate_water_percentage(binary_mask)
+            print(f"Water surface coverage: {water_percentage:.2f}%")
+
+        save_results(
+            image,
+            binary_mask,
+            overlay,
+            output_dir,
+            base_name,
+            save_mask=args.save_mask,
+            save_overlay=args.save_overlay
+        )
+
+        if args.save_results:
+            display_results(image, binary_mask, overlay, output_dir, base_name)
+
+        print("Inference completed successfully!")
+
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+"""
+Water Surface Segmentation Training Script
+Train YOLOv11n model for water surface segmentation on beach images.
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+from ultralytics import YOLO
+
+
+def parse_arguments() -> argparse.Namespace:
+    """Parse command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Train YOLOv11n model for water surface segmentation",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter
+    )
+
+    parser.add_argument(
+        "--data",
+        type=str,
+        required=True,
+        help="Path to data.yaml file"
+    )
+
+    parser.add_argument(
+        "--weights",
+        type=str,
+        default="yolov11n-seg.pt",
+        help="Path to pretrained weights"
+    )
+
+    parser.add_argument(
+        "--img",
+        type=int,
+        default=640,
+        help="Image size for training"
+    )
+
+    parser.add_argument(
+        "--batch",
+        type=int,
+        default=16,
+        help="Batch size"
+    )
+
+    parser.add_argument(
+        "--epochs",
+        type=int,
+        default=50,
+        help="Number of training epochs"
+    )
+
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="",
+        help="Device to use for training (cpu, cuda, mps)"
+    )
+
+    parser.add_argument(
+        "--project",
+        type=str,
+        default="runs/segment",
+        help="Project directory"
+    )
+
+    parser.add_argument(
+        "--name",
+        type=str,
+        default="nwsd_train",
+        help="Experiment name"
+    )
+
+    parser.add_argument(
+        "--patience",
+        type=int,
+        default=10,
+        help="Early stopping patience"
+    )
+
+    parser.add_argument(
+        "--save-period",
+        type=int,
+        default=5,
+        help="Save model every n epochs"
+    )
+
+    return parser.parse_args()
+
+
+def validate_inputs(args: argparse.Namespace) -> None:
+    """Validate input arguments."""
+    if not os.path.exists(args.data):
+        raise FileNotFoundError(f"Data configuration file not found: {args.data}")
+
+    if not args.weights.startswith("yolov11") and not os.path.exists(args.weights):
+        raise FileNotFoundError(f"Weights file not found: {args.weights}")
+
+
+def main():
+    """Main training function."""
+    args = parse_arguments()
+
+    try:
+        validate_inputs(args)
+
+        print(f"Loading model: {args.weights}")
+        model = YOLO(args.weights)
+
+        train_params = {
+            'data': args.data,
+            'imgsz': args.img,
+            'batch': args.batch,
+            'epochs': args.epochs,
+            'device': args.device,
+            'project': args.project,
+            'name': args.name,
+            'patience': args.patience,
+            'save_period': args.save_period,
+            'save': True,
+            'verbose': True,
+            'plots': True,
+            'val': True,
+        }
+
+        print("Starting training with parameters:")
+        for key, value in train_params.items():
+            print(f"  {key}: {value}")
+
+        results = model.train(**train_params)
+
+        model_save_path = os.path.join(args.project, args.name, "weights", "best.pt")
+        final_model_path = os.path.join("model", "nwsd-v2.pt")
+
+        os.makedirs("model", exist_ok=True)
+
+        if os.path.exists(model_save_path):
+            import shutil
+            shutil.copy2(model_save_path, final_model_path)
+            print(f"Best model saved to: {final_model_path}")
+
+        print("Training completed successfully!")
+
+    except Exception as e:
+        print(f"Error: {str(e)}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()