94 lines
3.2 KiB
Python
94 lines
3.2 KiB
Python
import rasterio
|
|
import numpy as np
|
|
from shapely.geometry import mapping
|
|
from rasterio.features import geometry_mask
|
|
|
|
# Step 1: Load Raster Data
|
|
def load_raster(raster_path):
|
|
with rasterio.open(raster_path) as src:
|
|
bands = [src.read(i) for i in range(1, src.count + 1)]
|
|
profile = src.profile
|
|
return np.stack(bands, axis=-1), profile
|
|
|
|
# Step 2: Extract Reference Spectra
|
|
def extract_reference_spectra(raster, training_shapefile, profile):
|
|
import geopandas as gpd
|
|
training_labels = gpd.read_file(training_shapefile)
|
|
|
|
reference_spectra = {}
|
|
|
|
for _, row in training_labels.iterrows():
|
|
class_id = row['class_id']
|
|
geom = [mapping(row['geometry'])]
|
|
|
|
mask = geometry_mask(
|
|
geom, transform=profile['transform'], invert=True, out_shape=(profile['height'], profile['width'])
|
|
)
|
|
|
|
masked_pixels = raster[mask]
|
|
mean_spectrum = np.mean(masked_pixels, axis=0)
|
|
reference_spectra[class_id] = mean_spectrum
|
|
|
|
return reference_spectra
|
|
|
|
# Step 3: Spectral Angle Mapping
|
|
def spectral_angle(pixel_spectrum, reference_spectrum):
|
|
# Avoid zero vectors
|
|
pixel_norm = np.linalg.norm(pixel_spectrum)
|
|
reference_norm = np.linalg.norm(reference_spectrum)
|
|
|
|
if pixel_norm == 0 or reference_norm == 0:
|
|
return np.pi # Maximum angle (180 degrees)
|
|
|
|
numerator = np.dot(pixel_spectrum, reference_spectrum)
|
|
denominator = pixel_norm * reference_norm
|
|
|
|
# Clip to avoid numerical issues with arccos
|
|
cosine_similarity = numerator / denominator
|
|
cosine_similarity = np.clip(cosine_similarity, -1.0, 1.0)
|
|
|
|
angle = np.arccos(cosine_similarity)
|
|
return angle
|
|
|
|
|
|
def classify_with_sam(raster, reference_spectra):
|
|
rows, cols, bands = raster.shape
|
|
classified_image = np.full((rows, cols), fill_value=255, dtype=np.uint8) # 255 as nodata
|
|
|
|
for i in range(rows):
|
|
for j in range(cols):
|
|
pixel_spectrum = raster[i, j, :]
|
|
angles = {class_id: spectral_angle(pixel_spectrum, ref_spectrum)
|
|
for class_id, ref_spectrum in reference_spectra.items()}
|
|
classified_image[i, j] = min(angles, key=angles.get) # Class with smallest angle
|
|
|
|
return classified_image
|
|
|
|
# Step 4: Save Classified Image
|
|
def save_classified_image(classified_image, profile, output_path):
|
|
profile.update(count=1, dtype=rasterio.uint8, nodata=255)
|
|
with rasterio.open(output_path, "w", **profile) as dst:
|
|
dst.write(classified_image, 1)
|
|
|
|
# Main Execution
|
|
if __name__ == "__main__":
|
|
# File paths
|
|
raster_path = "stacked/stacked.tif"
|
|
training_shapefile = "training/training.shp"
|
|
output_path = "classified/sam.tif"
|
|
|
|
# Load raster and metadata
|
|
raster, profile = load_raster(raster_path)
|
|
|
|
# Extract reference spectra
|
|
print("Extracting reference spectra...")
|
|
reference_spectra = extract_reference_spectra(raster, training_shapefile, profile)
|
|
|
|
# Classify raster using SAM
|
|
print("Classifying with SAM...")
|
|
classified_image = classify_with_sam(raster, reference_spectra)
|
|
|
|
# Save the classified image
|
|
save_classified_image(classified_image, profile, output_path)
|
|
print(f"Classified image saved to {output_path}")
|