segmentation.py 5.29 KB
Newer Older
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
#!/usr/bin/env python
# coding: utf-8

# SICOR is a freely available, platform-independent software designed to process hyperspectral remote sensing data,
# and particularly developed to handle data from the EnMAP sensor.

# This file contains image segmentation tools.

# Copyright (C) 2018  Niklas Bohn (GFZ, <nbohn@gfz-potsdam.de>),
# German Research Centre for Geosciences (GFZ, <https://www.gfz-potsdam.de>)

# This software was developed within the context of the EnMAP project supported by the DLR Space Administration with
# funds of the German Federal Ministry of Economic Affairs and Energy (on the basis of a decision by the German
# Bundestag: 50 EE 1529) and contributions from DLR, GFZ and OHB System AG.

# This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public
# License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later
# version.

# This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.

# You should have received a copy of the GNU General Public License along with this program.
# If not, see <https://www.gnu.org/licenses/>.


from scipy.linalg import eigh, norm
from scipy.spatial import KDTree
from scipy.stats import linregress
from skimage.segmentation import slic
31
import numpy as np
32
33
# from multiprocessing import Pool
from multiprocessing.pool import ThreadPool as Pool
34
35
36


def SLIC_segmentation(data_rad_all, n_pca, segs):
37
38
39
    """
    Perform superpixel segmentation by searching the image for a handful of geographically and compositionally
    representative locations. Create spatially-contiguous clusters of similar radiances.
40

41
    :param data_rad_all: radiance data array with shape (rows, cols, bands)
42
43
44
    :param n_pca:        number of principle components
    :param segs:         number of segments
    :return:             flattened radiance data array, number of segments, array of segment labels for each pixel
45
46
47
48
    """

    # Change from BIP to a long list of spectra
    rows, bands, cols = data_rad_all.shape[0], data_rad_all.shape[2], data_rad_all.shape[1]
49
    X = np.asarray(data_rad_all, dtype=np.float32).reshape([rows * cols, bands])
50
51
52

    # Excluding bad locations, calculate top PCA coefficients
    mu = X.mean(axis=0)
53
    C = np.cov(X, rowvar=False)
54
55
56
    [v, d] = eigh(C)

    # Determine segmentation compactness scaling based on top 5 eigenvalues
57
    cmpct = norm(np.sqrt(v[-n_pca:]))
58
59
60
61
62

    # Project, redimension as an image with 5 channels, and segment
    X_pca = (X - mu) @ d[:, -n_pca:]
    X_pca = X_pca.reshape([rows, cols, n_pca])

63
64
    labels = slic(X_pca, n_segments=segs, compactness=cmpct, max_iter=10, sigma=0, multichannel=True,
                  enforce_connectivity=True, min_size_factor=0.5, max_size_factor=3)
65
    lbl = np.unique(labels)
66
67
68
69
70
    segs = len(lbl)

    return X, segs, labels


71
72
73
def empirical_line_solution(X, rdn_subset, data_l2a_seg, rows, cols, bands, segs, labels, processes):
    """
    Apply empirical line solution to infer exact result of remaining spectra.
74
75
76

    :param X:            flattened radiance data array
    :param rdn_subset:   radiance data subset
77
    :param data_l2a_seg: segmented surface reflectance
78
79
80
81
82
    :param rows:         number of image rows
    :param cols:         number of image columns
    :param bands:        number of instrument bands
    :param segs:         number of segments
    :param labels:       array of segment labels for each pixel
83
    :param processes:    number of CPUs for multiprocessing
84
    :return:             extrapolated surface reflectance for each pixel
85
    """
86
87
88
89
90
91

    # First set up a matrix of locations, one per spectrum in both the reference set and the image cube.
    # We will simply use row and column indices for simplicity, rather than worrying about geographic or 3D locations
    row_grid, col_grid = np.meshgrid(np.arange(rows), np.arange(cols))
    locations = np.array([col_grid.flatten(), row_grid.flatten()]).T
    locations_subset = np.zeros((segs, 2))
92
93
94
95
96

    for i in range(segs):
        locations_subset[i, :] = locations[labels.flat == i, :].mean(axis=0)
    tree = KDTree(locations_subset)

97
98
99
100
101
102
    # compute slopes and intercepts in the same shape as radiance and reflectance
    unique_labels = np.unique(labels)
    global _globs
    _globs = dict(tree=tree, locs=locations_subset, k=15, rdn=rdn_subset, data=data_l2a_seg, bands=bands)
    with Pool(processes=processes) as pool:
        results = pool.map(_compute_coefficients_for_label, unique_labels)
103

104
105
106
    slopes, intercepts = np.empty([rows * cols, bands]), np.empty([rows * cols, bands])
    for lbl, coeffs in zip(unique_labels, results):
        slopes[labels.flat == lbl, :], intercepts[labels.flat == lbl, :] = coeffs
107

108
109
    # finally compute reflectance
    reflectance = X * slopes + intercepts
110

111
    return reflectance
112
113


114
_globs = dict()
115
116


117
118
119
120
121
def _compute_coefficients_for_label(label_idx):
    nn = _globs["tree"].query(_globs["locs"][label_idx, :], _globs["k"])[1]
    label_slopes, label_intercepts = zip(*[linregress(x=_globs["rdn"][:, nn, b],
                                                      y=_globs["data"][:, nn, b])[:2] for b in range(_globs["bands"])])
    return label_slopes, label_intercepts