Redundancy_predict_3c.py

#!/usr/bin/env python
# coding: utf-8
from PIL import Image
import cv2
from path import Path
from datasets2 import RSDataset
import torch
import torch.backends.cudnn
from torch.utils.data import DataLoader
import os
import tqdm
import numpy as np
import argparse
import shutil
import glob
from scipy import ndimage as ndi
from skimage.morphology import opening, closing, square
from pytorch_toolbelt.inference.tta import TTAWrapper, fliplr_image2mask, d4_image2mask
import matplotlib.pyplot


# need to create a file to store temp pictures
# if not os.path.exists('temp_pic'):
#     os.makedirs('temp_pic')
# path = './temp_pic/'
device = 'cuda:0'
dark = [0,0,0]

## use model to predict
def predict(model):
    result = []
    for images in tqdm.tqdm(test_loader):
        images = images.to(device, dtype=torch.float)
        temp = 0
        for keys in model.keys():
            model[keys].eval()
            net = TTAWrapper(model[keys], fliplr_image2mask)
            outputs = torch.sigmoid(model[keys](images))[0, 0, :, :]
            print(outputs.size())
            temp += outputs
        preds = temp/len(model)
        # preds = torch.from_numpy(preds)
        result.append(preds.detach().cpu().numpy())
    return result


def input_and_output(pic_path, model, generate_data):
    """
    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 = args.crop_size

    img = cv2.imread(pic_path)
    b = args.padding_size
    image = cv2.copyMakeBorder(img, b, b, b, b, cv2.BORDER_REFLECT)
    # height, width = ndsm.shape
    # ndsm = np.reshape(ndsm,(height,height,1))

    h, w = image.shape[0], image.shape[1]
    row = img.shape[0]//image_size
    col = img.shape[1]//image_size
    padding_img = np.zeros((h, w, 3), dtype=np.uint8)
    padding_img[0:h, 0:w, :] = image[:, :, :]

    padding_img = np.array(padding_img)
#     print ('src:',padding_img.shape)
    mask_whole = np.zeros((row*image_size, col*image_size), dtype=np.float32)
    if generate_data == False:
        result = predict(model)
        map_list = [str(i.name) for i in Path('temp_pic').files()]
    for i in range(row):
        for j in range(col):
            if generate_data:
                crop_img = redundancy_crop(padding_img, i, j, image_size)
                ch,cw,_ = crop_img.shape
                cv2.imwrite(f'temp_pic/{i}_{j}.png',crop_img)
            else:
                temp = result[map_list.index(f'{i}_{j}.png')]
                temp = redundancy_crop2(temp)
                mask_whole[i*image_size:i*image_size+image_size,j*image_size:j*image_size+image_size] = temp
    return mask_whole


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


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


def get_dataset_loaders( workers):
    batch_size = 1

    test_dataset = RSDataset(
        "./temp_pic"
    )
    test_loader = DataLoader(test_dataset, batch_size=batch_size, num_workers=workers)
    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


if __name__ =="__main__":
    from torch import nn
    from torch.nn import functional as F
    import torch
    import torchvision
    class Conv2dReLU(nn.Module):
        def __init__(self, in_channels, out_channels, kernel_size, padding=0,
                     stride=1, use_batchnorm=True, **batchnorm_params):

            super().__init__()

            layers = [
                nn.Conv2d(in_channels, out_channels, kernel_size,
                                  stride=stride, padding=padding, bias=not (use_batchnorm)),
                nn.ReLU(inplace=True),
            ]

            if use_batchnorm:
                layers.insert(1, nn.BatchNorm2d(out_channels, **batchnorm_params))

            self.block = nn.Sequential(*layers)

        def forward(self, x):
            return self.block(x)


    class DecoderBlock(nn.Module):
        def __init__(self, in_channels, out_channels,
                     use_batchnorm=True):
            super().__init__()
            self.block = nn.Sequential(
                Conv2dReLU(in_channels, out_channels, kernel_size=3, padding=1, use_batchnorm=use_batchnorm),
                Conv2dReLU(out_channels, out_channels, kernel_size=3, padding=1, use_batchnorm=use_batchnorm),
            )

        def forward(self, x):
            x, skip = x
            x = F.interpolate(x, scale_factor=2, mode='nearest')
            if skip is not None:
                x = torch.cat([x, skip], dim=1)
            x = self.block(x)
            return x


    class UNet(nn.Module):
        """
            UNet (https://arxiv.org/abs/1505.04597) with Resnet34(https://arxiv.org/abs/1512.03385) encoder

            """
        def __init__(self, num_classes=1,  pretrained=True, use_batchnorm=True, freeze_encoder=False):
            """
            :param num_classes:
            :param pretrained:
                False - no pre-trained network is used
                True  - encoder is pre-trained with resnet34
            :is_deconv:
                False: bilinear interpolation is used in decoder
                True: deconvolution is used in decoder
            """
            super().__init__()
            self.num_classes = num_classes
            self.pool = nn.MaxPool2d(2, 2)

            net = torchvision.models.resnet34(pretrained=pretrained)
            net.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)

            # with torch.no_grad():
            #     pretrained_conv1 = net.conv1.weight.clone()
            #     # Assign new conv layer with 4 input channels
            #     net.conv1 = torch.nn.Conv2d(4, 64, kernel_size=3, padding=1, bias=False)
            #     net.conv1.weight[:, :3] = pretrained_conv1
            #     net.conv1.weight[:, 3] = pretrained_conv1[:, 0]
            self.encoder = net

            decoder_channels = (256, 128, 64, 32, 16)
            encoder_channels = (512, 256, 128, 64, 64)
            in_channels = self.compute_channels(encoder_channels, decoder_channels)
            out_channels = decoder_channels

            for layer in self.encoder.parameters():
                layer.requires_grad = not freeze_encoder

            self.relu = nn.ReLU(inplace=True)

            self.conv0 = nn.Sequential(self.encoder.conv1,
                                       self.encoder.bn1,
                                       self.encoder.relu,
                                       self.pool)
            self.conv1 = self.encoder.layer1
            self.conv2 = self.encoder.layer2
            self.conv3 = self.encoder.layer3
            self.conv4 = self.encoder.layer4

            self.layer1 = DecoderBlock(in_channels[0], out_channels[0], use_batchnorm=use_batchnorm)
            self.layer2 = DecoderBlock(in_channels[1], out_channels[1], use_batchnorm=use_batchnorm)
            self.layer3 = DecoderBlock(in_channels[2], out_channels[2], use_batchnorm=use_batchnorm)
            self.layer4 = DecoderBlock(in_channels[3], out_channels[3], use_batchnorm=use_batchnorm)
            self.layer5 = DecoderBlock(in_channels[4], out_channels[4], use_batchnorm=use_batchnorm)
            self.final = nn.Conv2d(out_channels[4], num_classes, kernel_size=1)

        def compute_channels(self, encoder_channels, decoder_channels):
            channels = [
                encoder_channels[0] + encoder_channels[1],
                encoder_channels[2] + decoder_channels[0],
                encoder_channels[3] + decoder_channels[1],
                encoder_channels[4] + decoder_channels[2],
                0 + decoder_channels[3],
            ]
            return channels

        def forward(self, x):
            conv0 = self.encoder.conv1(x)
            conv0 = self.encoder.bn1(conv0)
            conv0 = self.encoder.relu(conv0)

            conv1 = self.pool(conv0)
            conv1 = self.conv1(conv1)
            conv2 = self.conv2(conv1)
            conv3 = self.conv3(conv2)
            conv4 = self.conv4(conv3)

            x = self.layer1([conv4, conv3])
            x = self.layer2([x, conv2])
            x = self.layer3([x, conv1])
            x = self.layer4([x, conv0])
            x = self.layer5([x, None])
            x = self.final(x)

            return x
# def my_predict():
    parse = argparse.ArgumentParser()
    parse.add_argument("--n_class", type=int, default=2, help="the number of classes")
    parse.add_argument("--model_name", type=str, default='UNet', help="UNet,PSPNet,FPN")

    parse.add_argument("--n_workers", type=int, default=4, help="the number of workers")
    parse.add_argument("--crop_size", type=int, default=256, help="the number of workers")
    parse.add_argument("--padding_size", type=int, default=32, help="the number of workers")
    parse.add_argument("--component_size", type=int, default=81, help="the size of component")

    args = parse.parse_args()
    # model_groups = ["UNet","PSPNet","FPN"]
    model_groups = ["model.pth"]
    # model_groups = ["FPN_val_0.85409586260551_epoch120.pth"]

# predict on more model
    models={}
    for index, item in enumerate(model_groups):
        models[item] = model = torch.load(f'./results_UNet/{item}', map_location='cuda:0')["model_state"]

    # model = torch.load(f'./results_{args.model_name}/{args.model_name}_weights_best.pth')["model_state"]
    imgList = glob.glob("./valid/*RGB.tif")
    num = len(imgList)

    save_path = f'./predict_YpUnet'
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for i in tqdm.tqdm(range(num)):
        if not os.path.exists('temp_pic'):
            os.makedirs('temp_pic')
        ### predict on one picture

        input_and_output(imgList[i], models, generate_data=True)
        name = os.path.split(imgList[i])[-1].split(".")[0]
        test_loader = get_dataset_loaders(args.n_workers)
        mask_result = input_and_output(imgList[i], models, generate_data=False)
        # 递归删除文件夹
        try:
            shutil.rmtree('temp_pic')
            shutil.rmtree('temp_ndsm')
        except:
            pass
        # dsm_ds = gdal.Open(imgList[i].replace("RGB", "DSM"), gdal.GA_ReadOnly)
        # band_dsm = dsm_ds.GetRasterBand(1)
        # nodata = band_dsm.GetNoDataValue()
        # dsm = band_dsm.ReadAsArray()
        #
        # mask_result[dsm == nodata] = 0
        threshold = 0.5
        mask_result[mask_result > threshold] = 1
        mask_result[mask_result <= threshold] = 0

        mask_result = opening(mask_result, square(6))
        labels = ndi.label(mask_result, output=np.uint32)[0]
        unique, counts = np.unique(labels, return_counts=True)
        for (k, v) in dict(zip(unique, counts)).items():
            if v < args.component_size:
                print("remove small objects successfully")
                mask_result[labels == k] = 0


        decoded = decode_segmap(mask_result, args.n_class)

        # print(mask_result.shape)
        cv2.imwrite(f'{save_path}/{name}.png', decoded)