utils/segmentation.py

"""
Segmentation utilities for image processing inspired by CLIPSeg techniques.
This is a simplified version that does not require the full transformers library.
"""

import os
import logging
import numpy as np
import cv2
from PIL import Image
from utils.geospatial import extract_contours, simplify_polygons, regularize_polygons, merge_nearby_polygons

def segment_by_color_threshold(image_path, output_path=None, 
                              threshold=127, color_channel=1, 
                              smoothing_sigma=1.0):
    """
    Segment an image based on color thresholding.
    This is a simple segmentation inspired by more complex models like CLIPSeg.
    
    Args:
        image_path (str): Path to the input image
        output_path (str, optional): Path to save the segmentation mask
        threshold (int): Pixel intensity threshold (0-255)
        color_channel (int): Color channel to use for thresholding (0=R, 1=G, 2=B)
        smoothing_sigma (float): Gaussian smoothing sigma
        
    Returns:
        numpy.ndarray: Segmentation mask
    """
    try:
        # Read the image
        img = cv2.imread(image_path)
        if img is None:
            # Try using PIL if OpenCV fails
            pil_img = Image.open(image_path).convert('RGB')
            img = np.array(pil_img)
            img = img[:, :, ::-1]  # RGB to BGR for OpenCV compatibility
        
        # Split channels and use the specified channel for segmentation
        b, g, r = cv2.split(img)
        channels = [r, g, b]
        
        if 0 <= color_channel < 3:
            channel = channels[color_channel]
        else:
            # Use grayscale if invalid channel specified
            channel = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Apply Gaussian blur to reduce noise
        if smoothing_sigma > 0:
            channel = cv2.GaussianBlur(channel, (0, 0), smoothing_sigma)
        
        # Apply thresholding to create binary mask
        _, mask = cv2.threshold(channel, threshold, 255, cv2.THRESH_BINARY)
        
        # Save the mask if output path is provided
        if output_path:
            cv2.imwrite(output_path, mask)
            logging.info(f"Saved segmentation mask to {output_path}")
            
        return mask
        
    except Exception as e:
        logging.error(f"Error in segmentation: {str(e)}")
        return None

def segment_by_adaptive_threshold(image_path, output_path=None, 
                                 block_size=11, c=2, 
                                 smoothing_sigma=1.0):
    """
    Segment an image using adaptive thresholding for better handling of
    lighting variations.
    
    Args:
        image_path (str): Path to the input image
        output_path (str, optional): Path to save the segmentation mask
        block_size (int): Size of the pixel neighborhood for threshold calculation
        c (int): Constant subtracted from the mean
        smoothing_sigma (float): Gaussian smoothing sigma
        
    Returns:
        numpy.ndarray: Segmentation mask
    """
    try:
        # Read the image
        img = cv2.imread(image_path)
        if img is None:
            # Try using PIL if OpenCV fails
            pil_img = Image.open(image_path).convert('RGB')
            img = np.array(pil_img)
            img = img[:, :, ::-1]  # RGB to BGR for OpenCV compatibility
        
        # Convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Apply Gaussian blur to reduce noise
        if smoothing_sigma > 0:
            gray = cv2.GaussianBlur(gray, (0, 0), smoothing_sigma)
        
        # Apply adaptive thresholding
        mask = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY, block_size, c
        )
        
        # Save the mask if output path is provided
        if output_path:
            cv2.imwrite(output_path, mask)
            logging.info(f"Saved segmentation mask to {output_path}")
            
        return mask
        
    except Exception as e:
        logging.error(f"Error in segmentation: {str(e)}")
        return None

def segment_by_otsu(image_path, output_path=None, smoothing_sigma=1.0):
    """
    Segment an image using Otsu's automatic thresholding method.
    
    Args:
        image_path (str): Path to the input image
        output_path (str, optional): Path to save the segmentation mask
        smoothing_sigma (float): Gaussian smoothing sigma
        
    Returns:
        numpy.ndarray: Segmentation mask
    """
    try:
        # Read the image
        img = cv2.imread(image_path)
        if img is None:
            # Try using PIL if OpenCV fails
            pil_img = Image.open(image_path).convert('RGB')
            img = np.array(pil_img)
            img = img[:, :, ::-1]  # RGB to BGR for OpenCV compatibility
        
        # Convert to grayscale
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Apply Gaussian blur to reduce noise
        if smoothing_sigma > 0:
            gray = cv2.GaussianBlur(gray, (0, 0), smoothing_sigma)
        
        # Apply Otsu's thresholding
        _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
        
        # Save the mask if output path is provided
        if output_path:
            cv2.imwrite(output_path, mask)
            logging.info(f"Saved segmentation mask to {output_path}")
            
        return mask
        
    except Exception as e:
        logging.error(f"Error in segmentation: {str(e)}")
        return None

def segment_and_extract_features(image_path, output_mask_path=None, 
                                feature_type="buildings", 
                                min_area=50, simplify_tolerance=2.0,
                                merge_distance=5.0):
    """
    Complete pipeline for segmentation and feature extraction.
    
    Args:
        image_path (str): Path to the input image
        output_mask_path (str, optional): Path to save the segmentation mask
        feature_type (str): Type of features to extract ("buildings", "trees", "water", "roads")
        min_area (int): Minimum feature area to keep
        simplify_tolerance (float): Tolerance for polygon simplification
        merge_distance (float): Distance for merging nearby polygons
        
    Returns:
        tuple: (mask, polygons) - Segmentation mask and list of simplified Shapely polygons
    """
    # Choose segmentation method based on feature type
    if feature_type.lower() == "buildings":
        # Buildings typically have clean edges and good contrast
        mask = segment_by_adaptive_threshold(
            image_path, output_mask_path, 
            block_size=15, c=2, smoothing_sigma=1.0
        )
    elif feature_type.lower() == "trees" or feature_type.lower() == "vegetation":
        # Trees typically strong in green channel
        mask = segment_by_color_threshold(
            image_path, output_mask_path,
            threshold=140, color_channel=1, smoothing_sigma=1.5
        )
    elif feature_type.lower() == "water":
        # Water typically has distinct spectral properties
        mask = segment_by_color_threshold(
            image_path, output_mask_path,
            threshold=120, color_channel=0, smoothing_sigma=2.0
        )
    else:
        # Default to Otsu for unknown feature types
        mask = segment_by_otsu(
            image_path, output_mask_path, smoothing_sigma=1.0
        )
    
    if mask is None:
        logging.error("Segmentation failed")
        return None, []
    
    # Save mask temporarily if needed for contour extraction
    temp_mask_path = None
    if not output_mask_path:
        temp_mask_path = os.path.join(
            os.path.dirname(image_path),
            f"{os.path.splitext(os.path.basename(image_path))[0]}_mask.png"
        )
        cv2.imwrite(temp_mask_path, mask)
        mask_path = temp_mask_path
    else:
        mask_path = output_mask_path
    
    # Extract contours from the mask
    polygons = extract_contours(mask_path, min_area=min_area)
    logging.info(f"Extracted {len(polygons)} initial polygons")
    
    # Clean up temporary file if created
    if temp_mask_path and os.path.exists(temp_mask_path):
        os.remove(temp_mask_path)
    
    # Simplify polygons
    polygons = simplify_polygons(polygons, tolerance=simplify_tolerance)
    
    # If buildings, regularize them to make more rectangular
    if feature_type.lower() == "buildings":
        polygons = regularize_polygons(polygons)
    
    # Merge nearby polygons to reduce count
    polygons = merge_nearby_polygons(polygons, distance_threshold=merge_distance)
    logging.info(f"After processing: {len(polygons)} polygons")
    
    return mask, polygons