# -*- ecoding: utf-8 -*-
# @ModuleName: pridict
# @Function: 
# @Author: <NAME>
# @Time: 2021/9/1 15:22
import os
import cv2
import keras
import numpy as np
import segmentation_models as sm
import albumentations as A
#DATA_DIR = 'C:/Users/ZhangYe/Desktop/core_segmentation0828/pr_imgs_2/'
dir = 'C:/Users/.../'
DATA_DIR = "C:/Users/..." + "/"
x_test_dir = os.path.join(DATA_DIR, 'test')
y_test_dir = os.path.join(DATA_DIR, 'testannot')
img_width = 640
img_height = 480
BACKBONE = 'efficientnetb5'
img_CLASSES = ['background', 'core']    # 图像标注几类，写几个
#checkpoint = "./best_core_model_res_adam2.h5"
checkpoint = "./model_eb5_adam.h5"
CLASSES = ['core']
#dir = "C:/Users/ZhangYe/Desktop/core_segmentation0828/data/CamVid" + "/"

BATCH_SIZE = 2

n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
activation = 'sigmoid' if n_classes == 1 else 'softmax'
model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)

class Dataset:
    """CamVid Dataset. Read images, apply augmentation and preprocessing transformations.

    Args:
        images_dir (str): path to images folder
        masks_dir (str): path to segmentation masks folder
        class_values (list): values of classes to extract from segmentation mask
        augmentation (albumentations.Compose): data transfromation pipeline
            (e.g. flip, scale, etc.)
        preprocessing (albumentations.Compose): data preprocessing
            (e.g. noralization, shape manipulation, etc.)

    """

    # CLASSES = ['sky', 'building', 'pole', 'road', 'pavement',
    #            'tree', 'signsymbol', 'fence', 'car',
    #            'pedestrian', 'bicyclist', 'unlabelled']

    CLASSES = img_CLASSES

    def __init__(
            self,
            images_dir,
            masks_dir,
            classes=None,
            augmentation=None,
            preprocessing=None,
    ):
        self.ids = os.listdir(images_dir)
        self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
        self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]

        # convert str names to class values on masks
        self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]

        self.augmentation = augmentation
        self.preprocessing = preprocessing

    def __getitem__(self, i):

        # read data
        image = cv2.imread(self.images_fps[i])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.masks_fps[i], 0)

        # extract certain classes from mask (e.g. cars)
        masks = [(mask == v) for v in self.class_values]
        mask = np.stack(masks, axis=-1).astype('float')

        # add background if mask is not binary
        if mask.shape[-1] != 1:
            background = 1 - mask.sum(axis=-1, keepdims=True)
            mask = np.concatenate((mask, background), axis=-1)

        # apply augmentations
        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        # apply preprocessing
        if self.preprocessing:
            sample = self.preprocessing(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        return image, mask

    def __len__(self):
        return len(self.ids)


class Dataloder(keras.utils.Sequence):
    """Load data from dataset and form batches

    Args:
        dataset: instance of Dataset class for image loading and preprocessing.
        batch_size: Integet number of images in batch.
        shuffle: Boolean, if `True` shuffle image indexes each epoch.
    """

    def __init__(self, dataset, batch_size=1, shuffle=False):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.indexes = np.arange(len(dataset))

        self.on_epoch_end()

    def __getitem__(self, i):

        # collect batch data
        start = i * self.batch_size
        stop = (i + 1) * self.batch_size
        data = []
        for j in range(start, stop):
            data.append(self.dataset[j])

        # transpose list of lists
        batch = [np.stack(samples, axis=0) for samples in zip(*data)]

        return batch

    def __len__(self):
        """Denotes the number of batches per epoch"""
        return len(self.indexes) // self.batch_size

    def on_epoch_end(self):
        """Callback function to shuffle indexes each epoch"""
        if self.shuffle:
            self.indexes = np.random.permutation(self.indexes)
def get_preprocessing(preprocessing_fn):
    """Construct preprocessing transform

    Args:
        preprocessing_fn (callbale): data normalization function
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose

    """

    _transform = [
        A.Lambda(image=preprocessing_fn),
    ]
    return A.Compose(_transform)


preprocess_input = sm.get_preprocessing(BACKBONE)
test_dataset = Dataset(
    x_test_dir,
    y_test_dir,
    classes=CLASSES,
    augmentation=None,
#     augmentation=get_validation_augmentation(),
    preprocessing=get_preprocessing(preprocess_input),
)

test_dataloader = Dataloder(test_dataset, batch_size=1, shuffle=False)

# check shapes for errors
images_size =  (BATCH_SIZE, img_height, img_width, 3)
masks_size =  (BATCH_SIZE, img_height, img_width, n_classes)
# assert test_dataloader[0][0].shape == images_size
# assert test_dataloader[0][1].shape == masks_size

# load best weights
model.load_weights(checkpoint)

# scores = model.evaluate_generator(test_dataloader)

# metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]

# print("Loss: {:.5}".format(scores[0]))
# for metric, value in zip(metrics, scores[1:]):
#     print("mean {}: {:.5}".format(metric.__name__, value))

n = 1
# ids = np.random.choice(np.arange(len(test_dataset)), size=n)
print(len(test_dataset),0)

for i in range(len(test_dataset)):
    image, mask = test_dataset[i] # get some sample:
    image = np.expand_dims(image, axis=0)
    pr_mask = model.predict(image)
    image = pr_mask.squeeze()
    # dir = "pr_imgs" + "\\"
    filename = dir + str(i) + ".png"
    print(filename)
    cv2.imwrite(filename, image)

# -*- ecoding: utf-8 -*-
# @ModuleName: pridict
# @Function: 
# @Author: <NAME>
# @Time: 2021/9/1 15:22
import os
import cv2
import keras
import numpy as np
import segmentation_models as sm
import albumentations as A
#DATA_DIR = 'C:/Users/ZhangYe/Desktop/core_segmentation0828/pr_imgs_2/'
dir = 'C:/Users/.../'
DATA_DIR = "C:/Users/..." + "/"
x_test_dir = os.path.join(DATA_DIR, 'test')
y_test_dir = os.path.join(DATA_DIR, 'testannot')
img_width = 640
img_height = 480
BACKBONE = 'efficientnetb5'
img_CLASSES = ['background', 'core']    # 图像标注几类，写几个
#checkpoint = "./best_core_model_res_adam2.h5"
checkpoint = "./model_eb5_adam.h5"
CLASSES = ['core']
#dir = "C:/Users/ZhangYe/Desktop/core_segmentation0828/data/CamVid" + "/"

BATCH_SIZE = 2

n_classes = 1 if len(CLASSES) == 1 else (len(CLASSES) + 1)  # case for binary and multiclass segmentation
activation = 'sigmoid' if n_classes == 1 else 'softmax'
model = sm.Unet(BACKBONE, classes=n_classes, activation=activation)

class Dataset:
    """CamVid Dataset. Read images, apply augmentation and preprocessing transformations.

    Args:
        images_dir (str): path to images folder
        masks_dir (str): path to segmentation masks folder
        class_values (list): values of classes to extract from segmentation mask
        augmentation (albumentations.Compose): data transfromation pipeline
            (e.g. flip, scale, etc.)
        preprocessing (albumentations.Compose): data preprocessing
            (e.g. noralization, shape manipulation, etc.)

    """

    # CLASSES = ['sky', 'building', 'pole', 'road', 'pavement',
    #            'tree', 'signsymbol', 'fence', 'car',
    #            'pedestrian', 'bicyclist', 'unlabelled']

    CLASSES = img_CLASSES

    def __init__(
            self,
            images_dir,
            masks_dir,
            classes=None,
            augmentation=None,
            preprocessing=None,
    ):
        self.ids = os.listdir(images_dir)
        self.images_fps = [os.path.join(images_dir, image_id) for image_id in self.ids]
        self.masks_fps = [os.path.join(masks_dir, image_id) for image_id in self.ids]

        # convert str names to class values on masks
        self.class_values = [self.CLASSES.index(cls.lower()) for cls in classes]

        self.augmentation = augmentation
        self.preprocessing = preprocessing

    def __getitem__(self, i):

        # read data
        image = cv2.imread(self.images_fps[i])
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        mask = cv2.imread(self.masks_fps[i], 0)

        # extract certain classes from mask (e.g. cars)
        masks = [(mask == v) for v in self.class_values]
        mask = np.stack(masks, axis=-1).astype('float')

        # add background if mask is not binary
        if mask.shape[-1] != 1:
            background = 1 - mask.sum(axis=-1, keepdims=True)
            mask = np.concatenate((mask, background), axis=-1)

        # apply augmentations
        if self.augmentation:
            sample = self.augmentation(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        # apply preprocessing
        if self.preprocessing:
            sample = self.preprocessing(image=image, mask=mask)
            image, mask = sample['image'], sample['mask']

        return image, mask

    def __len__(self):
        return len(self.ids)


class Dataloder(keras.utils.Sequence):
    """Load data from dataset and form batches

    Args:
        dataset: instance of Dataset class for image loading and preprocessing.
        batch_size: Integet number of images in batch.
        shuffle: Boolean, if `True` shuffle image indexes each epoch.
    """

    def __init__(self, dataset, batch_size=1, shuffle=False):
        self.dataset = dataset
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.indexes = np.arange(len(dataset))

        self.on_epoch_end()

    def __getitem__(self, i):

        # collect batch data
        start = i * self.batch_size
        stop = (i + 1) * self.batch_size
        data = []
        for j in range(start, stop):
            data.append(self.dataset[j])

        # transpose list of lists
        batch = [np.stack(samples, axis=0) for samples in zip(*data)]

        return batch

    def __len__(self):
        """Denotes the number of batches per epoch"""
        return len(self.indexes) // self.batch_size

    def on_epoch_end(self):
        """Callback function to shuffle indexes each epoch"""
        if self.shuffle:
            self.indexes = np.random.permutation(self.indexes)
def get_preprocessing(preprocessing_fn):
    """Construct preprocessing transform

    Args:
        preprocessing_fn (callbale): data normalization function
            (can be specific for each pretrained neural network)
    Return:
        transform: albumentations.Compose

    """

    _transform = [
        A.Lambda(image=preprocessing_fn),
    ]
    return A.Compose(_transform)


preprocess_input = sm.get_preprocessing(BACKBONE)
test_dataset = Dataset(
    x_test_dir,
    y_test_dir,
    classes=CLASSES,
    augmentation=None,
#     augmentation=get_validation_augmentation(),
    preprocessing=get_preprocessing(preprocess_input),
)

test_dataloader = Dataloder(test_dataset, batch_size=1, shuffle=False)

# check shapes for errors
images_size =  (BATCH_SIZE, img_height, img_width, 3)
masks_size =  (BATCH_SIZE, img_height, img_width, n_classes)
# assert test_dataloader[0][0].shape == images_size
# assert test_dataloader[0][1].shape == masks_size

# load best weights
model.load_weights(checkpoint)

# scores = model.evaluate_generator(test_dataloader)

# metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)]

# print("Loss: {:.5}".format(scores[0]))
# for metric, value in zip(metrics, scores[1:]):
#     print("mean {}: {:.5}".format(metric.__name__, value))

n = 1
# ids = np.random.choice(np.arange(len(test_dataset)), size=n)
print(len(test_dataset),0)

for i in range(len(test_dataset)):
    image, mask = test_dataset[i] # get some sample:
    image = np.expand_dims(image, axis=0)
    pr_mask = model.predict(image)
    image = pr_mask.squeeze()
    # dir = "pr_imgs" + "\\"
    filename = dir + str(i) + ".png"
    print(filename)
    cv2.imwrite(filename, image)

# -*- ecoding: utf-8 -*- # @ModuleName: pridict # @Function: # @Author: <NAME> # @Time: 2021/9/1 15:22 #DATA_DIR = 'C:/Users/ZhangYe/Desktop/core_segmentation0828/pr_imgs_2/' # 图像标注几类，写几个 #checkpoint = "./best_core_model_res_adam2.h5" #dir = "C:/Users/ZhangYe/Desktop/core_segmentation0828/data/CamVid" + "/" # case for binary and multiclass segmentation CamVid Dataset. Read images, apply augmentation and preprocessing transformations.
 
     Args:
         images_dir (str): path to images folder
         masks_dir (str): path to segmentation masks folder
         class_values (list): values of classes to extract from segmentation mask
         augmentation (albumentations.Compose): data transfromation pipeline
             (e.g. flip, scale, etc.)
         preprocessing (albumentations.Compose): data preprocessing
             (e.g. noralization, shape manipulation, etc.) # CLASSES = ['sky', 'building', 'pole', 'road', 'pavement', #            'tree', 'signsymbol', 'fence', 'car', #            'pedestrian', 'bicyclist', 'unlabelled'] # convert str names to class values on masks # read data # extract certain classes from mask (e.g. cars) # add background if mask is not binary # apply augmentations # apply preprocessing Load data from dataset and form batches
 
     Args:
         dataset: instance of Dataset class for image loading and preprocessing.
         batch_size: Integet number of images in batch.
         shuffle: Boolean, if `True` shuffle image indexes each epoch. # collect batch data # transpose list of lists Denotes the number of batches per epoch Callback function to shuffle indexes each epoch Construct preprocessing transform
 
     Args:
         preprocessing_fn (callbale): data normalization function
             (can be specific for each pretrained neural network)
     Return:
         transform: albumentations.Compose #     augmentation=get_validation_augmentation(), # check shapes for errors # assert test_dataloader[0][0].shape == images_size # assert test_dataloader[0][1].shape == masks_size # load best weights # scores = model.evaluate_generator(test_dataloader) # metrics = [sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5)] # print("Loss: {:.5}".format(scores[0])) # for metric, value in zip(metrics, scores[1:]): #     print("mean {}: {:.5}".format(metric.__name__, value)) # ids = np.random.choice(np.arange(len(test_dataset)), size=n) # get some sample: # dir = "pr_imgs" + "\\"