File size: 9,068 Bytes
747451d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
# /*---------------------------------------------------------------------------------------------
#  * 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 necessary libraries
import os
import sys
from pathlib import Path
import warnings
import sklearn
import mlflow
from hydra.core.hydra_config import HydraConfig
from omegaconf import DictConfig
from typing import Tuple, Optional, List, Dict
import numpy as np

# Suppress warnings and TensorFlow logs
warnings.filterwarnings("ignore")
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import onnxruntime
import tensorflow as tf
import tqdm

# Import utility functions
from image_classification.tf.src.preprocessing import postprocess_output, preprocess_input
from image_classification.tf.src.utils import ai_runner_invoke
from common.utils import (
    ai_runner_interp, ai_interp_input_quant, ai_interp_outputs_dequant,
    plot_confusion_matrix, log_to_file, display_figures
)  # Common utilities for evaluation and visualization


# Define a class for evaluating TFLite quantized models
class TFLiteQuantizedModelEvaluator:
    """

    A class to evaluate TensorFlow Lite (TFLite) quantized models.



    Args:

        cfg (DictConfig): Configuration object for evaluation.

        model (object): The quantized TFLite model to evaluate.

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

        Retrieves the evaluation target from the configuration.

        """
        if self.cfg.evaluation and self.cfg.evaluation.target:
            return self.cfg.evaluation.target
        return "host"

    def _get_interpreter(self, target):
        """

        Retrieves the AI runner interpreter for the specified target.

        Args: 

            target (str): target on which we intend to evaluate

        Returns:

            ai runner interpreter correctly parametrized

        """
        name_model = os.path.basename(self.quantized_model.model_path)
        return ai_runner_interp(target, name_model)

    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 model using {self.name_ds}...')
        target = self._get_target()     # Get the evaluation target
        ai_runner_interpreter = self._get_interpreter(target=target)    # Get the AI runner interpreter
        interpreter_quant = self.quantized_model    # Quantized TFLite model
        input_details = interpreter_quant.get_input_details()[0]
        input_index_quant = input_details["index"]
        output_index_quant = interpreter_quant.get_output_details()[0]["index"]
        output_details = interpreter_quant.get_output_details()[0]
        predictions_all = []    # Placeholder for all predictions
        test_pred = []      # Placeholder for predicted labels
        test_labels = []    # Placeholder for ground truth labels
        images_full = []    # Placeholder for processed input images

        # Iterate over the evaluation dataset
        for images, labels in tqdm.tqdm(self.eval_ds, total=len(self.eval_ds)):
            for image, label in zip(images, labels):
                # Preprocess the input image
                image_processed = preprocess_input(image, input_details)
                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(image_processed)
                
                # Perform inferences
                if target == 'host':
                    interpreter_quant.set_tensor(input_index_quant, image_processed)
                    interpreter_quant.invoke()
                    test_pred_score = interpreter_quant.get_tensor(output_index_quant)
                elif target in ['stedgeai_host', 'stedgeai_n6', 'stedgeai_h7p']:
                    image_preproc = ai_interp_input_quant(ai_runner_interpreter, image[None].numpy(), '.tflite')
                    test_pred_score = ai_runner_invoke(image_preproc, ai_runner_interpreter)
                    test_pred_score = ai_interp_outputs_dequant(ai_runner_interpreter, [test_pred_score])[0]
                    test_pred_score = np.reshape(test_pred_score, [1, -1])

                # Save predictions if configured
                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(test_pred_score)

                # Postprocess the output and store predictions
                predicted_label = postprocess_output(test_pred_score, output_details)
                test_pred.append(predicted_label)
                test_labels.append(label.numpy())

        # Save evaluation dataset in a .npy file if configured
        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:
                npy_in_name = getattr(self.cfg.evaluation, "npy_in_name", "unknown_npy_in_name")
                images_full = np.concatenate(images_full, axis=0)
                print("[INFO] : Shape of npy input dataset = {}".format(images_full.shape))
                np.save(os.path.join(self.output_dir, f"{npy_in_name}.npy"), images_full)

        # Save model output in a .npy file if configured
        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:
                npy_out_name = getattr(self.cfg.evaluation, "npy_out_name", "unknown_npy_out_name")
                predictions_all = np.concatenate(predictions_all, axis=0)
                print("[INFO] : Shape of npy predicted scores = {}".format(predictions_all.shape))
                np.save(os.path.join(self.output_dir, f"{npy_out_name}.npy"), predictions_all)

        # Compute the confusion matrix and accuracy
        labels = np.array(test_labels)
        logits = np.concatenate(test_pred, axis=0)
        logits = np.squeeze(logits)
        cm = sklearn.metrics.confusion_matrix(labels, logits)
        accuracy = round((np.sum(labels == logits) * 100) / len(test_labels), 2)

        # Log evaluation results
        print(f"[INFO] : Accuracy of quantized model on {self.name_ds} = {accuracy}%")
        log_to_file(self.output_dir,  f"Quantized model {self.name_ds}:")
        log_to_file(self.output_dir, f"Accuracy of quantized model : {accuracy} %")
        mlflow.log_metric(f"int_acc_{self.name_ds}", accuracy)

        # Plot and display the confusion matrix if enabled
        if self.display_figures:
            model_name = f"quantized_model_confusion_matrix_{self.name_ds}"
            plot_confusion_matrix(cm, class_names=self.class_names, model_name=model_name,
                                  title=f'{model_name}\naccuracy: {accuracy}', output_dir=self.output_dir)
            display_figures(self.cfg)

        return accuracy

    def evaluate(self):
        """

        Executes the full evaluation process.



        Returns:

            float: Accuracy of the quantized model on the evaluation dataset.

        """
        self._prepare_evaluation()      # Prepare the evaluation process
        acc = self._run_evaluate()      # Run the evaluation
        print('[INFO] : Evaluation complete.')
        return acc  # Return the accuracy