predict.py

#!/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()