Redundancy_predict_edge_only.py

#!/usr/bin/env python
# coding: utf-8
from PIL import Image
import cv2
from path import Path
from dataset import TestDataset
import torch
import torch.backends.cudnn
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch.nn.functional as F
import os
import gdalTools
import numpy as np
import shutil
import glob
import tqdm
from scipy import ndimage as ndi
from config_eval import ConfigEval
from skimage.morphology import remove_small_holes, closing, square, dilation, remove_small_objects, watershed
from skimage import io
from shutil import rmtree
import warnings
from osgeo import gdal
warnings.filterwarnings("ignore")

# need to create a file to store temp pictures
path = './temp_pic/'

device = 'cuda:0'
dark = [0, 0, 0]


# 水平翻转
def flip_horizontal_tensor(batch):
    columns = batch.data.size()[-1]
    return batch.index_select(-1, torch.LongTensor(list(reversed(range(columns)))).to(device))

#   垂直翻转
def flip_vertical_tensor(batch):
    rows = batch.data.size()[-2]
    return batch.index_select(-2, torch.LongTensor(list(reversed(range(rows)))).to(device))

# 水平翻转
def flip_horizontal_tensor_cpu(batch):
    columns = batch.data.size()[-1]
    return batch.index_select(-1, torch.LongTensor(list(reversed(range(columns)))))

#   垂直翻转
def flip_vertical_tensor_cpu(batch):
    rows = batch.data.size()[-2]
    return batch.index_select(-2, torch.LongTensor(list(reversed(range(rows)))))



def image_normalization(img, img_min=0, img_max=255,
                        epsilon=1e-12):
    """This is a typical image normalization function
    where the minimum and maximum of the image is needed
    source: https://en.wikipedia.org/wiki/Normalization_(image_processing)

    :param img: an image could be gray scale or color
    :param img_min:  for default is 0
    :param img_max: for default is 255

    :return: a normalized image, if max is 255 the dtype is uint8
    """

    img = np.float32(img)
    # whenever an inconsistent image
    img = (img - np.min(img)) * (img_max - img_min) / \
        ((np.max(img) - np.min(img)) + epsilon) + img_min
    return img


def input_and_output(pic_path, model, cfg, loader=None, generate_data=True, scale_size=None):
    """
    args:
        pic_path : the picture you want to predict
        model    : the model you want to predict
    note:
        step one : generate some pictures from one picture
        step two : predict from the images generated by step one 
    """
    image_size = cfg.crop_size

    data = gdal.Open(pic_path)
    lastChannel = data.RasterCount + 1
    arr = [data.GetRasterBand(idx).ReadAsArray() for idx in range(1, 4)]
    data = np.dstack(arr)

    raw_h, raw_w = data.shape[:2]

    # data = cv2.bilateralFilter(data, 9, 75, 75)

    b = cfg.padding_size
    row = raw_h // image_size + 1
    col = raw_w // image_size + 1
    radius_h = row * image_size - raw_h
    radius_w = col * image_size - raw_w
    data = cv2.copyMakeBorder(data, 0, radius_h, 0, radius_w, cv2.BORDER_REFLECT)

    data = cv2.copyMakeBorder(data, b, b, b, b, cv2.BORDER_REFLECT)

    h, w = data.shape[:2]

    # padding_img = data[:, :, :]

    data = np.array(data)
    mask_whole = np.zeros((row*image_size, col*image_size), dtype=np.uint8)
    if generate_data == False:
        print('starting prediction')
        result = []
        for batch in tqdm.tqdm(loader):
            # images = batch['img'].to(device, dtype=torch.float)
            images = Variable(batch['img'].to(device, dtype=torch.float))
            h, w = images.shape[2:]
            new_h = int(h * scale_size)
            images = F.interpolate(images, size=(new_h, new_h), mode='bilinear')
            temp = 0
            for keys in model.keys():
                # model[keys].eval()
                net = model[keys]
                net.eval()
                if cfg.TTA:
                    outputs1 = net(images)
                    outputs1 = [output1.cpu().detach() for output1 in outputs1]
                    outputs2 = net(flip_horizontal_tensor(images))
                    outputs2 = [output2.cpu().detach() for output2 in outputs2]

                    # outputs3 = net(flip_vertical_tensor(images))

                    outputs2 = [flip_horizontal_tensor_cpu(output) for output in outputs2]
                    # outputs3 = [flip_vertical_tensor(output) for output in outputs3]

                    outputs = [(output1 + output2)/3 for output1, output2 in zip(outputs1, outputs2)]
                    # outputs = [(output1 + output2 + output3)/4 for output1, output2, output3
                    #            in zip(outputs1, outputs2, outputs3)]

                    # outputs = [((output1 - torch.min(output1))/(torch.max(output1) - torch.min(output1)) +
                    #             (output2 - torch.min(output2))/(torch.max(output2) - torch.min(output2)))/2
                    #            for output1, output2 in zip(outputs1, outputs2)]

                else:
                    outputs = net(images)
                # tensor = outputs[-1][0, ...]
                # tensor = torch.mean(torch.stack(outputs, 0), 0)[0, ...]

                edge_maps = []
                for i in outputs:
                    tmp = torch.sigmoid(i).numpy()
                    edge_maps.append(tmp)
                tensor = np.array(edge_maps)

                preds = []
                for i in range(0, len(tensor)):
                    tmp_img = tensor[i][0, ...]
                    # tmp_img = tensor[i][0, 1, ...]
                    tmp_img = np.uint8(image_normalization(tmp_img)).squeeze()
                    tmp_img = cv2.bitwise_not(tmp_img)
                    # tmp_img[tmp_img < 0.0] = 0.0
                    # tmp_img = 255.0 * (1.0 - tmp_img)
                    # if i in [1, 2, 4]:
                    #     preds.append(tmp_img)
                    if i in [2, 3, 5]:
                        preds.append(tmp_img)
                    if i == 4:
                        fuse = tmp_img
                        # preds.append(tmp_img)

                fuse = fuse.astype(np.uint8)

                # Get the mean prediction of all the 7 outputs
                average = np.array(preds, dtype=np.float32)
                average = np.uint8(np.mean(average, axis=0))

                temp += average
            preds = temp / len(model)
            preds = cv2.resize(preds, (h, w))
            # preds = torch.from_numpy(preds)
            result.append(preds)
        map_list = [str(i.name) for i in Path('temp_pic').files()]
    for i in tqdm.tqdm(range(row)):
        for j in range(col):
            if generate_data:
                crop_img = redundancy_crop(data, i, j, image_size, cfg)
                ch,cw,_ = crop_img.shape
                io.imsave(f'temp_pic/{i}_{j}.png', crop_img.astype(np.uint8))
                # cv2.imwrite(f'temp_pic/{i}_{j}.tif', crop_img)
            else:
                temp = result[map_list.index(f'{i}_{j}.png')]
                temp = redundancy_crop2(temp, cfg)
                mask_whole[i*image_size:i*image_size+image_size, j*image_size:j*image_size+image_size] = temp
    return mask_whole[:raw_h, :raw_w]


def redundancy_crop(img, i, j, targetSize, cfg):
    if len(img.shape)>2:
        temp_img = img[i*targetSize:i*targetSize+targetSize+2*cfg.padding_size, j*targetSize:j*targetSize+targetSize+2*cfg.padding_size, :]
    else:
        temp_img = img[i*targetSize:i*targetSize+targetSize+2*cfg.padding_size, j*targetSize:j*targetSize+targetSize+2*cfg.padding_size]
    return temp_img


def redundancy_crop2(img, cfg):
    h = img.shape[0]
    w = img.shape[1]
    temp_img = img[cfg.padding_size:h-cfg.padding_size, cfg.padding_size:w-cfg.padding_size]
    return temp_img


def get_dataset_loaders():
    batch_size = 1

    test_dataset = TestDataset(
        "./temp_pic",
        mode='test'
    )

    test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=0)
    return test_loader


def get_labels():
    """Load the mapping that associates pascal classes with label colors

    Returns:
        np.ndarray with dimensions (2, 3)
    """
    return np.asarray(
        [
            [0, 0, 0],
            [255, 255, 255]
        ]
    )


def decode_segmap(label_mask, n_classes):
    """Decode segmentation class labels into a color image

    Args:
        label_mask (np.ndarray): an (M,N) array of integer values denoting
          the class label at each spatial location.
        plot (bool, optional): whether to show the resulting color image
          in a figure.

    Returns:
        (np.ndarray, optional): the resulting decoded color image.
    """
    label_colours = get_labels()
    r = label_mask.copy()
    g = label_mask.copy()
    b = label_mask.copy()
    for ll in range(0, n_classes):
        r[label_mask == ll] = label_colours[ll, 0]
        g[label_mask == ll] = label_colours[ll, 1]
        b[label_mask == ll] = label_colours[ll, 2]
    rgb = np.zeros((label_mask.shape[0], label_mask.shape[1], 3))
    rgb[:, :, 0] = r
    rgb[:, :, 1] = g
    rgb[:, :, 2] = b
    return rgb


###mask1   big    mask2 small
def my_watershed(mask1, mask2):
    """
    watershed from mask1 with markers from mask2
    """
    markers = ndi.label(mask2, output=np.uint32)[0]
    labels = watershed(mask1, markers, mask=mask1, watershed_line=True)
    return labels


def mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)


def predict(cfg):
    # 递归删除文件夹
    try:
        shutil.rmtree('temp_pic')
    except:
        pass

    if os.path.exists(cfg.save_path):
        rmtree(cfg.save_path)

    mkdir(cfg.save_path)
    # model_groups = glob.glob(cfg.model_path+'/*.pth')
    model_groups = [cfg.model_path]
    # imgList = glob.glob(cfg.data_path + '/*.tif')
    imgList = [cfg.data_path]
    num = len(imgList)

    # predict on more model
    print('loading models')
    models = {}
    for index, item in enumerate(model_groups):
        models[item] = torch.load(item, map_location=device)
        # models[item].load_segnet(r'D:\2021\3\EESNet_test\model_results\epoch253_model.pth', device)

    # for index, item in enumerate(model_groups):
    #     model = U2NET()
    #     model_dict = torch.load(item).module.state_dict()
    #     model.module.load_state_dict(model_dict)
    #     models[item] = model.to(device)


    for i in tqdm.tqdm(range(num)):
        im_proj, im_geotrans, im_width, im_height, im_data = gdalTools.read_img(imgList[i])
        if not os.path.exists('temp_pic'):
            os.makedirs('temp_pic')

        input_and_output(imgList[i], models, cfg, generate_data=True)

        # for j, scale_size in enumerate([1.125, 1, 0.875]):
        for j, scale_size in enumerate([1]):

            name = os.path.split(imgList[i])[-1].split(".")[0]
            test_loader = get_dataset_loaders()
            mask_result = input_and_output(imgList[i], models, cfg, loader=test_loader,
                                           generate_data=False,  scale_size=scale_size)
            mask_result = mask_result.astype(np.uint8)

            gdalTools.write_img(os.path.join(cfg.save_path, name + f'_{j}.tif'), im_proj, im_geotrans, 255 - mask_result)

        # # 递归删除文件夹
        # try:
        #     shutil.rmtree('temp_pic')
        # except:
        #     pass

    fuse_root = cfg.save_path + '_ms'
    if os.path.exists(fuse_root):
        rmtree(fuse_root)

    mkdir(fuse_root)
    for i in tqdm.tqdm(range(num)):
        im_proj, im_geotrans, im_width, im_height, im_data = gdalTools.read_img(imgList[i])
        baseName = os.path.basename(imgList[i]).split('.')[0]
        multiList = glob.glob(f'{cfg.save_path}/{baseName}*.tif')
        imgs = []
        for imgPath in multiList:
            im_proj, im_geotrans, im_width, im_height, im_data = gdalTools.read_img(imgPath)
            im_data = (im_data - np.min(im_data)) / (np.max(im_data) - np.min(im_data)) * 255
            imgs.append(im_data.astype(np.uint8))

        imgs = np.mean(np.stack(imgs, 0), 0)
        gdalTools.write_img(f'{fuse_root}/{baseName}.tif', im_proj, im_geotrans, imgs.astype(np.uint8))
        imgs2 = image_sharpening(imgs.copy()).astype(np.uint8)

        # Otsu's thresholding after Gaussian filtering
        blur = cv2.GaussianBlur(imgs2, (5, 5), 0)
        # ret3, th3 = cv2.threshold(blur, 0, 255, cv2.THRESH_OTSU)
        # imgs2 = remove_small_objects(th3, min_size=100)
        # imgs2 = remove_small_holes(imgs2, area_threshold=30)

        gdalTools.write_img(f'{fuse_root}/{baseName}_sharpening.tif', im_proj, im_geotrans, blur.astype(np.uint8))

        # forestpath = imgList[i].replace('test', 'test_forest')
        # forest = cv2.imread(forestpath, 0)
        # ret1, th1 = cv2.threshold(forest, 0, 255, cv2.THRESH_OTSU)
        # forest_edge = cv2.Canny(th1, 30, 100)
        # forest_edge = dilation(forest_edge, square(5))
        # imgs3 = np.where(forest_edge > 0, imgs2 + 40, imgs2)
        # gdalTools.write_img(f'{fuse_root}/{baseName}_sharpening_forest.tif', im_proj, im_geotrans, imgs3.astype(np.uint8))


def image_sharpening(image):
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]], np.float32)  # 定义一个核
    dst = cv2.filter2D(image, -1, kernel=kernel)
    return dst


if __name__ == '__main__':
    cfg = ConfigEval()
    f = open('config_test.txt')
    data = f.readlines()
    cfg.data_path = data[0].replace('\n', '')
    cfg.model_path = data[1].replace('\n', '')
    outRoot = data[2].replace('\n', '')

    predict(cfg)

    # mkdir('predict37')
    # imglist = glob.glob('predict36_a/*.tif')
    #
    # imgs = []
    # for imgPath in imglist:
    #     im_proj, im_geotrans, im_width, im_height, im_data = gdalTools.read_img(imgPath)
    #     imgs.append(im_data)
    #
    # imgs = np.mean(np.stack(imgs, 0), 0)
    # print(imgs.shape)
    # # cv2.imwrite('./predict37/cq_test_all.tif', imgs.astype())
    # gdalTools.write_img('./predict37/cq_test_all.tif', im_proj, im_geotrans, imgs.astype(np.uint8))