rscls.py

# -*- coding: utf-8 -*-
"""
This is a script for satellite image classification
Last updated on Aug 6 2019
@author: Shengjie Liu
@Email: liushengjie0756@gmail.com
@functions
1. generate samples from satellite images
2. grid search SVM/random forest parameters
3. object-based post-classification refinement
superpixel-based regularization for classification maps
4. confusion matrix: OA, kappa, PA, UA, AA
5. save maps as images
@sample codes
c = rscls.rscls(image,ground_truth,cls=number_of_classes)
c.padding(patch)
c.normalize(style='-11')  # optional
    
x_train,y_train = c.train_sample(num_per_cls)
x_train,y_train = rscls.make_sample(x_train,y_train)
    
x_test,y_test = c.test_sample()
# for superpixel refinement
c.locate_obj(seg)
pcmap = rscls.obpc(c.seg,predicted,c.obj)
@Notes
Ground truth file should be uint8 format begin with 1
Background = 0
"""



import numpy as np
import copy
import scipy.stats as stats
from sklearn.svm import SVC
from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier
from sklearn.naive_bayes import GaussianNB
import matplotlib.pyplot as plt

class rscls:
    def __init__(self,im,gt,cls):
        if cls==0:
            print('num of class not specified !!')
        self.im = copy.deepcopy(im)
        if gt.max()!=cls:
            self.gt = copy.deepcopy(gt-1)
        else:
            self.gt = copy.deepcopy(gt-1)
        self.gt_b = copy.deepcopy(gt)
        self.cls = cls
        self.patch = 1
        self.imx,self.imy,self.imz = self.im.shape
        self.record = []
        self.sample = {}
        
    def padding(self,patch):
        self.patch = patch
        pad = self.patch//2
        r1 = np.repeat([self.im[0,:,:]], pad, axis=0)
        r2 = np.repeat([self.im[-1,:,:]], pad, axis=0)
        self.im = np.concatenate((r1, self.im, r2))
        r1 = np.reshape(self.im[:,0,:],[self.imx + 2 * pad, 1, self.imz])
        r2 = np.reshape(self.im[:,-1,:],[self.imx + 2 * pad, 1, self.imz])
        r1 = np.repeat(r1, pad, axis=1)
        r2 = np.repeat(r2, pad, axis=1)
        self.im = np.concatenate((r1, self.im, r2), axis=1)
        self.im = self.im.astype('float32')

    def normalize(self,style='01'):
        im = self.im
        for i in range(im.shape[-1]):
            im[:,:,i]=(im[:,:,i]-im[:,:,i].min())/(im[:,:,i].max()-im[:,:,i].min())
        if style == '-11':
            im = im*2-1
    
    def locate_sample(self):
        sam = []
        for i in range(self.cls):
            _xy = np.array(np.where(self.gt==i)).T
            _sam = np.concatenate([_xy,i*np.ones([_xy.shape[0],1])],axis=-1)
            try:
                sam = np.concatenate([sam,_sam],axis=0)
            except:
                sam = _sam
        self.sample = sam.astype(int)
        
    def get_patch(self, xy):
        d = self.patch//2
        x = xy[0]
        y = xy[1]
        try:
            self.im[x][y]
        except IndexError:
            return []
        x += d
        y += d
        sam = self.im[(x - d):(x + d + 1), (y - d):(y + d + 1)]
        return np.array(sam)
    
    def train_sample(self,pn):
        x_train,y_train = [],[]
        self.locate_sample()
        _samp = self.sample
        for _cls in range(self.cls):
            _xy = _samp[_samp[:,2]==_cls]
            np.random.shuffle(_xy)
            _xy = _xy[:pn,:]
            for xy in _xy:
                self.gt[xy[0],xy[1]] = 255 # !!
                #
                x_train.append(self.get_patch(xy[:-1]))
                y_train.append(xy[-1])
            # print(_xy)
        x_train,y_train = np.array(x_train), np.array(y_train)
        idx = np.random.permutation(x_train.shape[0])
        x_train = x_train[idx]
        y_train = y_train[idx]
        return x_train,y_train.astype(int)
    
    def test_sample(self):
        x_test,y_test = [],[]
        self.locate_sample()
        _samp = self.sample
        for _cls in range(self.cls):
            _xy = _samp[_samp[:,2]==_cls]
            np.random.shuffle(_xy)
            for xy in _xy:
                x_test.append(self.get_patch(xy[:-1]))
                y_test.append(xy[-1])
        return np.array(x_test), np.array(y_test)
    
    def all_sample(self):
        imx,imy = self.gt.shape
        sample = []
        for i in range(imx):
            for j in range(imy):
                sample.append(self.get_patch(np.array([i,j])))
        return np.array(sample)
    
    def all_sample_light(self,clip=0,bs=10):
        imx,imy = self.gt.shape
        imz = self.im.shape[-1]
        patch = self.patch
        # fp = np.memmap('allsample' + str(clip) + '.h5', dtype='float32', mode='w+', shape=(imgx*self.IMGY,5,5,bs))
        fp = np.zeros([imx*imy,patch,patch,imz])
        countnum = 0
        for i in range(imx*clip,imx*(clip+1)):
            for j in range(imy):
                xy = np.array([i,j])
                fp[countnum,:,:,:] = self.get_patch(xy)
                countnum += 1
        return fp
    
    def all_sample_row_hd(self,sub=0):
        imx,imy = self.gt.shape
        imz = self.im.shape[-1]
        patch = self.patch
        # fp = np.memmap('allsample' + str(clip) + '.h5', dtype='float32', mode='w+', shape=(imgx*self.IMGY,5,5,bs))
        fp = np.zeros([imx*imy,patch,patch,imz])
        countnum = 0
        for i in range(sub):
            for j in range(imy):
                xy = np.array([i,j])
                fp[countnum,:,:,:] = self.get_patch(xy)
                countnum += 1
        return fp
    
    def all_sample_row(self,sub=0):
        imx,imy = self.gt.shape
        fp = []
        for j in range(imy):
            xy = np.array([sub,j])
            fp.append(self.get_patch(xy))
        return np.array(fp)

    def all_sample_heavy(self,name,clip=0,bs=10):
        imx,imy = self.gt.shape
        imz = self.im.shape[-1]
        patch = self.patch
        try:
            fp = np.memmap(name, dtype='float32', mode='w+', shape=(imx*imy,patch,patch,imz))
        except:
            fp = np.memmap(name, dtype='float32', mode='r', shape=(imx*imy,patch,patch,imz))
        # fp = np.zeros([imx*imy,patch,patch,imz])
        countnum = 0
        for i in range(imx*clip,imx*(clip+1)):
            for j in range(imy):
                xy = np.array([i,j])
                fp[countnum,:,:,:] = self.get_patch(xy)
                countnum += 1
        return fp
    
    def read_all_sample(self,name,clip=0,bs=10):
        imx,imy = self.gt.shape
        imz = self.im.shape[-1]
        patch = self.patch
        fp = np.memmap(name, dtype='float32', mode='r', shape=(imx*imy,patch,patch,imz))
        return fp
    
    def locate_obj(self,seg):
        obj = {}
        for i in range(seg.min(),seg.max()+1):
            obj[str(i)] = np.where(seg==i)
        self.obj = obj
        self.seg = seg
        
def obpc(seg,cmap,obj):
    pcmap = copy.deepcopy(cmap)
    for (k,v) in obj.items():
        tmplabel = stats.mode(cmap[v])[0]
        pcmap[v] = tmplabel
    return pcmap
           
def cfm(pre, ref, ncl=9):
    if ref.min() != 0:
        print('warning: label should begin with 0 !!')
        return

    nsize = ref.shape[0]
    cf = np.zeros((ncl,ncl))
    for i in range(nsize):
        cf[pre[i], ref[i]] += 1
    
    tmp1 = 0
    for j in range(ncl):
        tmp1 = tmp1 + (cf[j,:].sum()/nsize)*(cf[:,j].sum()/nsize)
    cfm = np.zeros((ncl+2,ncl+1))
    cfm[:-2,:-1] = cf
    oa = 0
    for i in range(ncl):
        if cf[i,:].sum():
            cfm[i,ncl] = cf[i,i]/cf[i,:].sum()
        if cf[:,i].sum():
            cfm[ncl,i] = cf[i,i]/cf[:,i].sum()
        oa += cf[i,i]
    cfm[-1, 0] = oa/nsize
    cfm[-1, 1] = (cfm[-1, 0]-tmp1)/(1-tmp1)
    cfm[-1, 2] = cfm[ncl,:-1].mean()
    print('oa: ', format(cfm[-1,0],'.5'), ' kappa: ', format(cfm[-1,1],'.5'),
          ' mean: ', format(cfm[-1,2],'.5'))
    return cfm

def gtcfm(pre,gt,ncl):
    if gt.max()==255:
        print('warning: max 255 !!')
    cf = np.zeros([ncl,ncl])
    for i in range(gt.shape[0]):
        for j in range(gt.shape[1]):
            if gt[i,j]:
                cf[pre[i,j]-1,gt[i,j]-1] += 1
    tmp1 = 0
    nsize = np.sum(gt!=0)
    for j in range(ncl):
        tmp1 = tmp1 + (cf[j,:].sum()/nsize)*(cf[:,j].sum()/nsize)
    cfm = np.zeros((ncl+2,ncl+1))
    cfm[:-2,:-1] = cf
    oa = 0
    for i in range(ncl):
        if cf[i,:].sum():
            cfm[i,ncl] = cf[i,i]/cf[i,:].sum()
        if cf[:,i].sum():
            cfm[ncl,i] = cf[i,i]/cf[:,i].sum()
        oa += cf[i,i]
    cfm[-1, 0] = oa/nsize
    cfm[-1, 1] = (cfm[-1, 0]-tmp1)/(1-tmp1)
    cfm[-1, 2] = cfm[ncl,:-1].mean()
    print('oa: ', format(cfm[-1,0],'.5'), ' kappa: ', format(cfm[-1,1],'.5'),
          ' mean: ', format(cfm[-1,2],'.5'))
    return cfm

def svm(trainx,trainy):
    cost = []
    gamma = []
    for i in range(-5,16,2):
        cost.append(np.power(2.0,i))
    for i in range(-15,4,2):
        gamma.append(np.power(2.0,i))
                
    parameters = {'C':cost,'gamma':gamma}
    svm = SVC(verbose=0,kernel='rbf')
    clf = GridSearchCV(svm, parameters,cv=3)
    p = clf.fit(trainx, trainy)
            
    print(clf.best_params_)
    bestc = clf.best_params_['C']
    bestg = clf.best_params_['gamma']
    tmpc = [-1.75,-1.5,-1.25,-1,-0.75,-0.5,-0.25,0.0,
            0.25,0.5,0.75,1.0,1.25,1.5,1.75]
    cost = []
    gamma=[]
    for i in tmpc:
        cost.append(bestc*np.power(2.0,i))
        gamma.append(bestg*np.power(2.0,i))
    parameters = {'C':cost,'gamma':gamma}
    svm = SVC(verbose=0,kernel='rbf')
    clf = GridSearchCV(svm, parameters,cv=3)
    p = clf.fit(trainx, trainy)
    print(clf.best_params_)
    p2 = clf.best_estimator_
    return p2

def svm_rbf(trainx,trainy):
    cost = []
    gamma = []
    for i in range(-3,10,2):
        cost.append(np.power(2.0,i))
    for i in range(-5,4,2):
        gamma.append(np.power(2.0,i))
                
    parameters = {'C':cost,'gamma':gamma}
    svm = SVC(verbose=0,kernel='rbf')
    clf = GridSearchCV(svm, parameters,cv=3)
    clf.fit(trainx, trainy)
            
    #print(clf.best_params_)
    bestc = clf.best_params_['C']
    bestg = clf.best_params_['gamma']
    tmpc = [-1.75,-1.5,-1.25,-1,-0.75,-0.5,-0.25,0.0,
            0.25,0.5,0.75,1.0,1.25,1.5,1.75]
    cost = []
    gamma=[]
    for i in tmpc:
        cost.append(bestc*np.power(2.0,i))
        gamma.append(bestg*np.power(2.0,i))
    parameters = {'C':cost,'gamma':gamma}
    svm = SVC(verbose=0,kernel='rbf')
    clf = GridSearchCV(svm, parameters,cv=3)
    clf.fit(trainx, trainy)
    #print(clf.best_params_)
    p = clf.best_estimator_
    return p

def rf(trainx,trainy,sim=1,nj=1):
    nest = []
    nfea = []
    for i in range(20, 201, 20):
        nest.append(i)
    if sim:
        for i in range(1,int(trainx.shape[-1])):
            nfea.append(i)
        parameters = {'n_estimators':nest,'max_features':nfea}
    else:
        parameters = {'n_estimators':nest}
    rf = RandomForestClassifier(n_jobs=nj,verbose=0,oob_score=False)
    clf = GridSearchCV(rf, parameters, cv=3)
    p = clf.fit(trainx, trainy)
    p2 = clf.best_estimator_
    return p2

def GNB(trainx,trainy):
    clf = GaussianNB()
    p = clf.fit(trainx, trainy)
    return p

def svm_linear(trainx,trainy):
    cost = []
    for i in range(-3,10,2):
        cost.append(np.power(2.0,i))
                
    parameters = {'C':cost}
    svm = SVC(verbose=0,kernel='linear')
    clf = GridSearchCV(svm, parameters,cv=3)
    clf.fit(trainx, trainy)
            
    #print(clf.best_params_)
    bestc = clf.best_params_['C']
    tmpc = [-1.75,-1.5,-1.25,-1,-0.75,-0.5,-0.25,0.0,
            0.25,0.5,0.75,1.0,1.25,1.5,1.75]
    cost = []
    for i in tmpc:
        cost.append(bestc*np.power(2.0,i))
    parameters = {'C':cost}
    svm = SVC(verbose=0,kernel='linear')
    clf = GridSearchCV(svm, parameters,cv=3)
    clf.fit(trainx, trainy)
    p = clf.best_estimator_
    return p

def make_sample(sample, label):
    a = np.flip(sample,1)
    b = np.flip(sample,2)
    c = np.flip(b,1)
    newsample = np.concatenate((a,b,c,sample),axis=0)
    newlabel = np.concatenate((label,label,label,label),axis=0)
    return newsample, newlabel

def save_cmap(img, cmap, fname):
   
    sizes = np.shape(img)
    height = float(sizes[0])
    width = float(sizes[1])
     
    fig = plt.figure()
    fig.set_size_inches(width/height, 1, forward=False)
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
 
    ax.imshow(img, cmap=cmap)
    plt.savefig(fname, dpi = height) 
    plt.close()