evaluation.py

# ================================================================
# MIT License
# Copyright (c) 2021 edwardyehuang (https://github.com/edwardyehuang)
# ================================================================

import tensorflow as tf

from tqdm import tqdm

from iseg.metrics.mean_iou import MeanIOU

from iseg.core_model import SegBase

from iseg.losses.catecrossentropy_ignore_label import catecrossentropy_ignore_label_loss
from iseg.metrics.seg_metric_wrapper import SegMetricWrapper

def evaluate(
    distribute_strategy,
    model,
    data,
    batch_size,
    num_class,
    ignore_label=255,
    scale_rates=[0.5, 0.75, 1.0, 1.25, 1.5, 1.75],
    flip=True,
    val_image_count=0,
    pre_compute_fns=[],
):

    if not isinstance(model, SegBase):
        raise ValueError("ALl model must based on SegBase")

    ds = prepare_dataset(distribute_strategy, data, batch_size, val_image_count=val_image_count)

    processed_count = 0

    print("Evaluate mIOU : ")

    with distribute_strategy.scope():

        loss_func = catecrossentropy_ignore_label_loss(
            num_class=num_class, ignore_label=ignore_label, batch_size=batch_size, reduction=False
        )

        loss_metrics = tf.keras.metrics.Mean("loss")

        iou_metrics = MeanIOU(num_class)

        iou_metrics = SegMetricWrapper(iou_metrics, num_class=num_class, ignore_label=ignore_label, name="IOU")

        if pre_compute_fns is not None and isinstance(pre_compute_fns, list):
            for fn in pre_compute_fns:
                iou_metrics.add_pre_compute_fn(fn)

        print("Current image format = {}".format(tf.keras.backend.image_data_format()))

        with tqdm(total=val_image_count) as pbar:
            for inputs in ds:
                eval_step(inputs, model, scale_rates, flip, loss_func, loss_metrics, [iou_metrics], distribute_strategy)

                processed_count += batch_size

                pbar.update(batch_size)
                pbar.set_description(
                    "Processed : {:}, current loss = {:4f}, current IOU = {:.2f} %".format(
                        processed_count, loss_metrics.result(), iou_metrics.result().numpy() * 100
                    )
                )

    mean_loss = loss_metrics.result()
    mean_iou = iou_metrics.result()
    
    print("-----------------------------------------------")

    print(f"Mean loss on val set : {mean_loss.numpy()}")
    print(f"Mean IoU on val set : {mean_iou.numpy()}")

    print("-----------------------------------------------")

    print(f"Per-class IoU on val set :")

    pre_class_iou =  iou_metrics.metric.per_class_result()

    print(pre_class_iou)

    loss_metrics.reset_states()
    iou_metrics.reset_states()

    return mean_iou


@tf.function(autograph=False)
def eval_step(ds_inputs, model: SegBase, scale_rates, flip, loss_func, loss_metrics, mertics, distribute_strategy):
    def step_fn(inputs):
        images, labels = inputs

        predictions = model.inference_with_multi_scales(images, training=False, scale_rates=scale_rates, flip=flip)

        # predictions = model.inference(images, False)

        loss = loss_func(labels, predictions)

        loss_metrics.update_state(loss)

        for mertic in mertics:
            mertic.update_state(labels, predictions)

    return distribute_strategy.run(step_fn, args=(ds_inputs,))


def prepare_dataset(distribute_strategy, data, batch_size=16, val_image_count=0):

    # AUTOTUNE = tf.data.experimental.AUTOTUNE

    ds = data
    # ds = ds.repeat()
    ds = ds.batch(batch_size, drop_remainder=False)

    options = tf.data.Options()
    options.experimental_distribute.auto_shard_policy = tf.data.experimental.AutoShardPolicy.DATA
    ds = ds.with_options(options)

    # ds = ds.prefetch(buffer_size=AUTOTUNE)

    ds = distribute_strategy.experimental_distribute_dataset(ds)

    return ds