40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
import argparse
|
|
import rasterio
|
|
import numpy as np
|
|
from sklearn.cluster import KMeans
|
|
|
|
def main(input_file, output_file, n_clusters):
|
|
with rasterio.open(input_file) as src:
|
|
raster_data = src.read()
|
|
profile = src.profile
|
|
|
|
rows, cols = raster_data.shape[1], raster_data.shape[2]
|
|
raster_data_reshaped = raster_data.reshape(raster_data.shape[0], -1).T
|
|
|
|
kmeans = KMeans(n_clusters=n_clusters, random_state=0)
|
|
kmeans.fit(raster_data_reshaped)
|
|
cluster_labels = kmeans.labels_
|
|
|
|
classified_raster = cluster_labels.reshape(rows, cols)
|
|
|
|
with rasterio.open(
|
|
output_file, 'w',
|
|
driver='GTiff',
|
|
height=rows,
|
|
width=cols,
|
|
count=1,
|
|
dtype=classified_raster.dtype,
|
|
crs=src.crs,
|
|
transform=src.transform,
|
|
) as dst:
|
|
dst.write(classified_raster, 1)
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Unsupervised classification for raster data.")
|
|
parser.add_argument("input_file", type=str, help="Path to the input .tif file.")
|
|
parser.add_argument("output_file", type=str, help="Path to save the classified .tif file.")
|
|
parser.add_argument("--n_clusters", type=int, default=4, help="Number of clusters for classification (default: 4).")
|
|
args = parser.parse_args()
|
|
|
|
main(args.input_file, args.output_file, args.n_clusters)
|