object_detection/tf/src/evaluation/tflite_evaluator.py

# /*---------------------------------------------------------------------------------------------
#  * Copyright (c) 2022 STMicroelectronics.
#  * All rights reserved.
#  *
#  * This software is licensed under terms that can be found in the LICENSE file in
#  * the root directory of this software component.
#  * If no LICENSE file comes with this software, it is provided AS-IS.
#  *--------------------------------------------------------------------------------------------*/

import os
import warnings
import numpy as np
import tqdm
import mlflow
import shutil
import tensorflow as tf
import matplotlib.pyplot as plt
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig
from timeit import default_timer as timer
from datetime import timedelta
from tabulate import tabulate
from pathlib import Path

from object_detection.tf.src.postprocessing import get_detections
from object_detection.tf.src.models import model_family
from object_detection.tf.src.utils import ai_runner_invoke
from object_detection.tf.src.utils import ObjectDetectionMetricsData, calculate_objdet_metrics, calculate_average_metrics
from common.utils import (
    ai_runner_interp, ai_interp_input_quant, ai_interp_outputs_dequant, log_to_file, display_figures
)  # Common utilities for evaluation and visualization


class TFLiteQuantizedModelEvaluator:
    """
    A class to evaluate TensorFlow Lite (TFLite) quantized object detection models.

    Args:
        cfg (DictConfig): Configuration object for evaluation.
        model (object): The quantized TFLite Interpreter object.
        dataloaders (dict): Dictionary containing datasets for testing and validation.
    """
    def __init__(self, cfg: DictConfig, model: object,
                 dataloaders: dict = None):
        self.cfg = cfg
        self.quantized_model = model
        self.test_ds = dataloaders['test']
        self.valid_ds = dataloaders['valid']
        self.output_dir = HydraConfig.get().runtime.output_dir
        self.class_names = cfg.dataset.class_names
        self.display_figures = cfg.general.display_figures
        self.eval_ds = None
        self.name_ds = None
    
    def _prepare_evaluation(self):
        """
        Prepares the evaluation process by selecting the appropriate dataset.
        """
        # Use the test dataset if available; otherwise, use the validation dataset
        if self.test_ds:
            self.eval_ds = self.test_ds
            self.name_ds = "test_set"
        else:
            self.eval_ds = self.valid_ds
            self.name_ds = "validation_set"

    def _get_target(self):
        if self.cfg.evaluation and self.cfg.evaluation.target:
            return self.cfg.evaluation.target
        return "host"

    def _get_interpreter(self, target):
        name_model = os.path.basename(self.quantized_model.model_path)
        return ai_runner_interp(target, name_model)
    
    def _display_objdet_metrics(self, metrics, class_names):
        table = []
        classes = list(metrics.keys())    
        for c in sorted(classes):
            table.append([
                class_names[c],
                round(100 * metrics[c].pre, 1),
                round(100 * metrics[c].rec, 1),
                round(100 * metrics[c].ap, 1)])
                
        print()
        headers = ["Class name", "Precision %", "  Recall %", "   AP %  "]
        print()
        print(tabulate(table, headers=headers, tablefmt="pipe", numalign="center"))

        mpre, mrec, mAP = calculate_average_metrics(metrics)
        
        print("\nAverages over classes %:")
        print("-----------------------")
        print(" Mean precision: {:.1f}".format(100 * mpre))
        print(" Mean recall:    {:.1f}".format(100 * mrec))
        print(" Mean AP (mAP):  {:.1f}".format(100 * mAP))

    def _plot_precision_versus_recall(self, metrics, class_names, plots_dir):
        """
        Plot the precision versus recall curves. AP values are the areas under these curves.
        """

        # Create the directory where plots will be saved
        if os.path.exists(plots_dir):
            shutil.rmtree(plots_dir)
        os.makedirs(plots_dir)

        for c in list(metrics.keys()):
            
            # Plot the precision versus recall curve
            figure = plt.figure(figsize=(10, 10))
            plt.xlabel("recall")
            plt.ylabel("interpolated precision")
            plt.title("Class '{}' (AP = {:.2f})".
                        format(class_names[c], metrics[c].ap * 100))
            plt.plot(metrics[c].interpolated_precision, metrics[c].interpolated_recall)
            plt.grid()

            # Save the plot in the plots directory
            plt.savefig(f"{plots_dir}/{class_names[c]}.png")
            plt.close(figure)

    def _run_evaluate(self):
        """
        Runs the evaluation process and computes metrics.

        Returns:
            float: Accuracy of the quantized model on the evaluation dataset.
        """
        tf.print(f'[INFO] : Evaluating the quantized object detection model using {self.name_ds}...')
        target = self._get_target()
        ai_runner_interpreter = self._get_interpreter(target=target)
        input_details = self.quantized_model.get_input_details()[0]
        output_details = self.quantized_model.get_output_details()
        model_batch_size = input_details['shape_signature'][0]
        if model_batch_size != 1 and target == 'host':
            batch_size = 64
        else:
            batch_size = 1

        input_shape = tuple(input_details['shape'][1:])
        image_size = input_shape[:2]
        dataset_size = sum([x.shape[0] for x, _ in self.eval_ds])

        exmpl, _ = iter(self.eval_ds).next()
        batch_size = exmpl.shape[0]

        _, labels = iter(self.eval_ds).next()
        num_labels = int(tf.shape(labels)[1])

        cpp = self.cfg.postprocessing
        metrics_data = None
        num_detections = 0
        predictions_all = []
        images_full = []

        start_time = timer()

        for i, data in enumerate(tqdm.tqdm(self.eval_ds)):
            images, gt_labels = data
            batch_size = int(tf.shape(images)[0])

            input_index = input_details['index']
            tensor_shape = (batch_size,) + input_shape
            self.quantized_model.resize_tensor_input(input_index, tensor_shape)
            self.quantized_model.allocate_tensors()

            scale, zero_point = input_details['quantization']
            images_quant = images / scale + zero_point
            input_dtype = input_details['dtype']
            images_quant = tf.cast(images_quant, input_dtype)
            images_quant = tf.clip_by_value(images_quant, np.iinfo(input_dtype).min, np.iinfo(input_dtype).max)
            
            if "evaluation" in self.cfg and self.cfg.evaluation:
                if "gen_npy_input" in self.cfg.evaluation and self.cfg.evaluation.gen_npy_input==True: 
                    images_full.append(images)

            if target == 'host':
                self.quantized_model.set_tensor(input_index, images_quant)
                self.quantized_model.invoke()
            elif target in ['stedgeai_host', 'stedgeai_n6', 'stedgeai_h7p']:
                data_quant = ai_interp_input_quant(ai_runner_interpreter, images.numpy(), '.tflite')
                predictions = ai_runner_invoke(data_quant, ai_runner_interpreter)
                predictions = ai_interp_outputs_dequant(ai_runner_interpreter, predictions)
            else:
                raise RuntimeError(f"Unsupported target: {target}")
            
            if model_family(self.cfg.model.model_type) in ["ssd", "st_yoloxn"]:
                if target == 'host':
                    # Model outputs are scores, boxes and anchors.
                    predictions = (self.quantized_model.get_tensor(output_details[0]['index']),
                                   self.quantized_model.get_tensor(output_details[1]['index']),
                                   self.quantized_model.get_tensor(output_details[2]['index']))
            else:
                if target == 'host':
                    predictions = self.quantized_model.get_tensor(output_details[0]['index'])
                elif target in ['stedgeai_host', 'stedgeai_n6', 'stedgeai_h7p']:
                    predictions = predictions[0]
                    
            if "evaluation" in self.cfg and self.cfg.evaluation:
                if "gen_npy_output" in self.cfg.evaluation and self.cfg.evaluation.gen_npy_output==True:
                    predictions_all.append(predictions)

            # The TFLITE version of yolov8 has channel-first outputs
            if model_family(self.cfg.model.model_type) in ["yolov8n"]:
                predictions = tf.transpose(predictions, perm=[0, 2, 1])

            boxes, scores = get_detections(self.cfg, predictions, image_size)

            if i == 0:
                num_detections = boxes.shape[1]
                metrics_data = ObjectDetectionMetricsData(
                    num_labels, cpp.max_detection_boxes, len(self.class_names),
                    num_detections, dataset_size, batch_size
                )

            metrics_data.add_data(gt_labels, boxes, scores)
            metrics_data.update_batch_index(i, cpp.confidence_thresh, cpp.NMS_thresh, image_size)
        
        # Saves evaluation dataset in a .npy
        if "evaluation" in self.cfg and self.cfg.evaluation:
            if "gen_npy_input" in self.cfg.evaluation and self.cfg.evaluation.gen_npy_input==True: 
                if "npy_in_name" in self.cfg.evaluation and self.cfg.evaluation.npy_in_name:
                    npy_in_name = self.cfg.evaluation.npy_in_name
                else:
                    npy_in_name = "unknown_npy_in_name"
                images_full = np.concatenate(images_full, axis=0)
                np.save(os.path.join(self.output_dir, f"{npy_in_name}.npy"), images_full)

        # Saves model output in a .npy
        if "evaluation" in self.cfg and self.cfg.evaluation:
            if "gen_npy_output" in self.cfg.evaluation and self.cfg.evaluation.gen_npy_output==True: 
                if "npy_out_name" in self.cfg.evaluation and self.cfg.evaluation.npy_out_name:
                    npy_out_name = self.cfg.evaluation.npy_out_name
                else:
                    npy_out_name = "unknown_npy_out_name"
                predictions_all = np.concatenate(predictions_all, axis=0)
                np.save(os.path.join(self.output_dir, f"{npy_out_name}.npy"), predictions_all)

        end_time = timer()
        eval_run_time = int(end_time - start_time)
        print("Evaluation run time: " + str(timedelta(seconds=eval_run_time)))

        groundtruths, detections = metrics_data.get_data()
        metrics = calculate_objdet_metrics(groundtruths, detections, cpp.IoU_eval_thresh)

        self._display_objdet_metrics(metrics, self.class_names)

        log_to_file(self.output_dir, f"Quantized TFLite object detection model dataset used: {self.cfg.dataset.dataset_name}")

        mpre, mrec, mAP = calculate_average_metrics(metrics)
        model_type = "quantized"
        log_to_file(self.output_dir, "{}_model_mpre: {:.1f}".format(model_type, 100 * mpre))
        log_to_file(self.output_dir, "{}_model_mrec: {:.1f}".format(model_type, 100 * mrec))
        log_to_file(self.output_dir, "{}_model_map: {:.1f}".format(model_type, 100 * mAP))
        
        # Log metrics in mlflow
        mlflow.log_metric(f"{model_type}_model_mpre", round(100 * mpre, 2))
        mlflow.log_metric(f"{model_type}_model_mrec", round(100 * mrec, 2))
        mlflow.log_metric(f"{model_type}_model_mAP", round(100 * mAP, 2))
        
        if self.cfg.postprocessing.plot_metrics:
            print("\nPlotting precision versus recall curves")
            plots_dir = os.path.join(self.output_dir, "precision_vs_recall_curves", os.path.basename(self.quantized_model.model_path))
            print("Plots directory:", plots_dir)
            self._plot_precision_versus_recall(metrics, self.class_names, plots_dir)

        print('[INFO] : Evaluation complete.')
        return metrics

    def evaluate(self):
        """
        Executes the full evaluation process.

        Returns:
            dict: Dictionary of evaluation metrics for each class.
        """
        self._prepare_evaluation()      # Prepare the evaluation process
        return self._run_evaluate()