Spaces:
Sleeping
Sleeping
File size: 8,811 Bytes
5cf313a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 |
"""
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 |