Spaces:
Sleeping
Sleeping
| """ | |
| 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 |