Redundancy_predict_segmentation.py

#!/usr/bin/env python
# coding: utf-8
import cv2
from path import Path
from dataset import SasDataset
import torch
import torch.backends.cudnn
from torch.utils.data import DataLoader
from torch.autograd import Variable
import torch.nn.functional as F

import gdalTools
import numpy as np
import shutil
import glob
import tqdm
import os
from scipy import ndimage as ndi
from config_eval import ConfigEval
from skimage.morphology import remove_small_holes, closing, square, opening, 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)
    mask_whole_dn = np.zeros((row*image_size, col*image_size), dtype=np.uint8)

    if generate_data == False:
        print('starting prediction')
        result = []
        result_dn = []
        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
            temp_dn = 0
            for keys in model.keys():
                # model[keys].eval()
                net = model[keys]
                net.eval()
                if cfg.TTA:
                    outputs1 = net(images, seg=True).cpu().detach()
                    outputs2 = flip_horizontal_tensor(net(flip_horizontal_tensor(images), seg=True)).cpu().detach()

                    # outputs1 = net(images).cpu().detach()
                    # outputs2 = flip_horizontal_tensor(net(flip_horizontal_tensor(images))).cpu().detach()

                    outputs = (outputs2 + outputs1)/2.

                else:
                    outputs = net(images, seg=True)

                prob = F.softmax(outputs, dim=1)
                # prob = outputs
                prob = prob[0]
                _, maxprob = torch.max(prob, 0)

                fuse = maxprob.numpy()

                prob = torch.sigmoid(outputs).numpy()
                tmp_img = prob[0, 1, ...]
                tmp_img = np.uint8(image_normalization(tmp_img)).squeeze()
                tmp_img = cv2.bitwise_not(tmp_img)

                temp += fuse
                temp_dn += tmp_img
            preds = temp / len(model)
            preds = cv2.resize(preds, (h, w))
            # preds = torch.from_numpy(preds)
            result.append(preds)
            result_dn.append(temp_dn)
        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)
                temp_dn = result_dn[map_list.index(f'{i}_{j}.png')]
                temp_dn = redundancy_crop2(temp_dn, cfg)
                mask_whole[i*image_size:i*image_size+image_size, j*image_size:j*image_size+image_size] = temp
                mask_whole_dn[i*image_size:i*image_size+image_size, j*image_size:j*image_size+image_size] = temp_dn
    return mask_whole[:raw_h, :raw_w], mask_whole_dn[: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 = SasDataset(
        "./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_seg(cfg):
    save_path = cfg.save_path + '_seg'
    if os.path.exists(save_path):
        rmtree(save_path)

    mkdir(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)
        # models[item] = torch.load(item, map_location='cuda:0')["ema_state_dict"]

    # model = torch.load(f'./results_{args.model_name}/{args.model_name}_weights_best.pth')["model_state"]

    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]):

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

            # mask_result = np.where(mask_result > 0, 1, 0)

            gdalTools.write_img(os.path.join(save_path, name + f'.tif'), im_proj, im_geotrans, mask_result)
            # gdalTools.write_img(os.path.join(save_path, name + f'_dn.tif'), im_proj, im_geotrans, mask_result_dn)

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

    # fuse_root = cfg.save_path + '_ms'
    # 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))


if __name__ == '__main__':
    pwd = os.path.dirname(os.path.abspath(__file__))

    cfg = ConfigEval()
    f = open('config_test.txt')
    data = f.readlines()
    cfg.data_path = data[0].replace('\n', '')
    cfg.model_path = data[1].replace('\n', '')
    outPath = data[2].replace('\n', '')
    outRoot = os.path.split(outPath)[0]
    outName = os.path.split(outPath)[-1].split('.')[0]

    predict_seg(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.tif', im_proj, im_geotrans, imgs.astype(np.uint8))