-
Notifications
You must be signed in to change notification settings - Fork 0
/
image_classification.py
34 lines (24 loc) · 1.11 KB
/
image_classification.py
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
import numpy as np
# import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from functions import *
###########################################################
# Parameters setting #
###########################################################
class_num = 5
F_tb_path = r"data/Coleambally/S2_20211223.tif"
F_tb_class_path = rf"data/Coleambally/S2_20211223_KMeans-{class_num}.tif"
if __name__ == "__main__":
image, profile = read_raster(F_tb_path)
print(profile)
image_pct2 = linear_pct_stretch(image, 2)
image_pct2 = color_composite(image_pct2, [3, 2, 1])
kmeans = KMeans(n_clusters=class_num, max_iter=300)
kmeans.fit(image.reshape(-1, image.shape[2]))
F_class = kmeans.labels_.reshape(image.shape[0], image.shape[1])
# fig, axes = plt.subplots(1, 2, sharex=True, sharey=True)
# axes[0].imshow(image_pct2)
# axes[1].imshow(F_class, cmap="Paired")
# plt.show()
write_raster(np.expand_dims(F_class, axis=2), profile, F_tb_class_path)
print("Classified the fine image using K-Means algorithm!")