steveyin/stm32-modelzoo-app
0
1# /*---------------------------------------------------------------------------------------------
2# * Copyright (c) 2025 STMicroelectronics.
3# * All rights reserved.
4# *
5# * This software is licensed under terms that can be found in the LICENSE file in
6# * the root directory of this software component.
7# * If no LICENSE file comes with this software, it is provided AS-IS.
8# *--------------------------------------------------------------------------------------------*/
9import collections
10import fnmatch
11import texttable
12import os
13
14# for common registries
15from common.registries.dataset_registry import DATASET_WRAPPER_REGISTRY
16from common.registries.model_registry import MODEL_WRAPPER_REGISTRY
17from common.registries.trainer_registry import TRAINER_WRAPPER_REGISTRY
18from common.registries.quantizer_registry import QUANTIZER_WRAPPER_REGISTRY
19from common.registries.evaluator_registry import EVALUATOR_WRAPPER_REGISTRY
20from common.registries.predictor_registry import PREDICTOR_WRAPPER_REGISTRY
21from common.utils import LOGGER
22
23# for tensorflow based image classification
24import image_classification.tf.wrappers.datasets # this is done so that registeration happens before get_dataloader is called
25import image_classification.tf.wrappers.models # this is done so that registeration happens before get_model is called
26import image_classification.tf.wrappers.training # this is done so that registeration happens before get_trainer is called
27import image_classification.tf.wrappers.quantization # this is done so that registeration happens before get_quantizer is called
28import image_classification.tf.wrappers.evaluation # this is done so that registeration happens before get_evaluator is called
29import image_classification.tf.wrappers.prediction # this is done so that registeration happens before get_predictor is called
30
31# for pytorch based image classification
32import image_classification.pt.wrappers.datasets # this is done so that registeration happens before get_dataloader is called
33import image_classification.pt.wrappers.models # this is done so that registeration happens before get_model is called
34import image_classification.pt.wrappers.training # this is done so that registeration happens before get_trainer is called
35import image_classification.pt.wrappers.quantization # this is done so that registeration happens before get_quantizer is called
36import image_classification.pt.wrappers.evaluation # this is done so that registeration happens before get_evaluator is called
37import image_classification.pt.wrappers.prediction # this is done so that registeration happens before get_evaluator is called
38
39# for tensorflow based object detection
40import object_detection.tf.wrappers.models.standard_models
41import object_detection.tf.wrappers.models.custom_models
42import object_detection.tf.wrappers.datasets # this is done so that registeration happens before get_dataloader is called
43import object_detection.tf.wrappers.prediction # this is done so that registeration happens before get_predictor is called
44import object_detection.tf.wrappers.evaluation # this is done so that registeration happens before get_evaluator is called
45import object_detection.tf.wrappers.quantization # this is done so that registeration happens before get_quantizer is called
46import object_detection.tf.wrappers.training # this is done so that registeration happens before get_trainer is called
47
48# for pytorch based object detection
49import object_detection.pt.wrappers.models
50import object_detection.pt.wrappers.datasets
51import object_detection.pt.wrappers.training
52import object_detection.pt.wrappers.evaluation
53
54from common.model_utils.tf_model_loader import load_model_from_path
55from pathlib import Path
56import torch
57import tensorflow as tf
58import onnxruntime
59
60
61__all__ = ['get_dataloaders',
62 "get_model",
63 "get_trainer",
64 "get_quantizer",
65 "get_evaluator",
66 "get_predictor",
67 "list_models",
68 "list_models_by_dataset",]
69
70def get_dataloaders(cfg):
71 """
72 Tries to find a matching dataloader creation wrapper function in the registry and uses it to create a new dataloder dict.
73
74 returns datasplits in the following format:
75 {
76 'train': train_data_loader,
77 'valid': valid_data_loader,
78 'quantization': quantization_data_loader,
79 'test' : test_data_loader
80 'predict' : predict_data_loader
81 }
82 TODO : add other keys like val, quant
83 """
84 #if mode not in ["benchmarking", "deployment"]:
85 if cfg.dataset.dataset_name != "<unnamed>":
86 data_split_wrapper_fn = DATASET_WRAPPER_REGISTRY.get(framework=cfg.model.framework,
87 dataset_name=cfg.dataset.dataset_name,
88 use_case=cfg.use_case)
89 # The registered function expects a single cfg object
90 return data_split_wrapper_fn(cfg)
91 else:
92 return {'train': None, 'valid': None, 'quantization': None, 'test': None, 'predict': None,}
93
94
95def get_model(cfg):
96 """
97 Tries to find a matching model creation wrapper function in the registry and uses it to create a new model object.
98 """
99 allowed_exts = [".keras", ".h5", ".tflite", ".onnx"]
100 model_path = getattr(cfg.model, "model_path", None)
101 if model_path:
102 _, ext = os.path.splitext(model_path)
103 if ext.lower() in allowed_exts:
104 LOGGER.info(f"Loading model from {model_path}")
105 model = load_model_from_path(cfg, model_path)
106 return model
107 # LOGGER.info(f"Loading model from {model_path}")
108 # _, ext = os.path.splitext(model_path)
109 # if ext.lower() in allowed_exts:
110 # model = load_model_from_path(cfg, model_path)
111 #
112 # Covers cases where there is no model_path provided, or pt checkpoints
113
114 model_func = MODEL_WRAPPER_REGISTRY.get(model_name=cfg.model.model_name.lower(),
115 use_case=cfg.use_case,
116 framework=cfg.model.framework)
117
118 # The registered function expects a single cfg object
119 model = model_func(cfg)
120 saved_model_dir = os.path.join(cfg.output_dir, cfg.general.saved_models_dir)
121 os.makedirs(saved_model_dir, exist_ok=True)
122 if cfg.model.framework == 'tf':
123 saved_model = Path(saved_model_dir, cfg.model.model_name, cfg.model.model_name+".keras")
124 saved_model.parent.mkdir(exist_ok=True)
125 model.save(saved_model)
126 setattr(model, 'model_path', saved_model)
127 cfg.model.model_path = saved_model
128
129 return model
130
131
132def get_trainer(dataloaders, model, cfg):
133 """
134 Returns an instance of the trainer class from registry.
135 """
136 trainer_cls = TRAINER_WRAPPER_REGISTRY.get(
137 trainer_name=cfg.training.trainer_name,
138 framework=cfg.model.framework,
139 use_case=cfg.use_case
140 )
141
142 # The registered function expects dataloaders, model and full config
143 return trainer_cls(dataloaders=dataloaders, model=model, cfg=cfg)
144
145
146def get_quantizer(dataloaders, model, cfg):
147 """
148 Returns an instance of the quantizer class from registry.
149 """
150 quantizer_cls = QUANTIZER_WRAPPER_REGISTRY.get(
151 quantizer_name=cfg.quantization.quantizer.lower(),
152 framework=cfg.model.framework,
153 use_case=cfg.use_case
154 )
155
156 # The registered function expects dataloaders, model and full config
157 return quantizer_cls(dataloaders=dataloaders, model=model, cfg=cfg)
158
159def get_evaluator(dataloaders, model, cfg):
160 """
161 Returns an instance of the evaluator class from registry.
162 """
163 if isinstance(model, tf.keras.Model):
164 evaluator_name = "keras_evaluator"
165 elif 'Interpreter' in str(type(model)):
166 evaluator_name = "tflite_evaluator"
167 elif isinstance(model, onnxruntime.InferenceSession):
168 evaluator_name = "onnx_evaluator"
169 elif isinstance(model, torch.nn.Module):
170 model_name = cfg.model.model_name.lower()
171 # Will not work if user is providing custom format.
172 if "ssd" in model_name:
173 evaluator_name = "ssd"
174 # ---- YOLOD----
175 elif "yolod" in model_name:
176 evaluator_name = "yolod"
177 else :
178 evaluator_name = "torch_evaluator"
179 else:
180 raise TypeError("Unsupported model type for evaluation")
181
182 evaluator_cls = EVALUATOR_WRAPPER_REGISTRY.get(
183 evaluator_name=evaluator_name,
184 framework=cfg.model.framework,
185 use_case=cfg.use_case,
186 )
187
188 return evaluator_cls(dataloaders=dataloaders, model=model, cfg=cfg)
189
190
191def get_predictor(dataloaders, model, cfg):
192 """
193 Returns an instance of the predictor class from registry.
194 """
195 if isinstance(model, tf.keras.Model):
196 predictor_name = "keras_predictor"
197 elif 'Interpreter' in str(type(model)):
198 predictor_name = "tflite_predictor"
199 elif isinstance(model, onnxruntime.InferenceSession):
200 predictor_name = "onnx_predictor"
201 else:
202 raise TypeError("Unsupported model type for predictor")
203
204 predictor_cls = PREDICTOR_WRAPPER_REGISTRY.get(
205 predictor_name=predictor_name,
206 framework=cfg.model.framework,
207 use_case=cfg.use_case
208 )
209
210 # The registered function expects dataloaders, model and full config
211 return predictor_cls(dataloaders=dataloaders, model=model, cfg=cfg)
212
213
214def list_models(
215 filter_string='',
216 match_all=True,
217 print_table=True,
218 with_checkpoint=False,
219):
220 """
221 A helper function to list all existing models based on text filters
222 You can provide list of strings like model_name, use_case, framework.
223 It will print a table of corresponding available models
224
225 :param filter: a string or list of strings containing model name, use_case , framework or "model_name_use_case_framework"
226 to use as a filter
227 :param print_table: Whether to print a table with matched models (if False, return as a list)
228 """
229 #print(MODEL_WRAPPER_REGISTRY.registry_dict.keys())
230 if with_checkpoint:
231 all_model_keys = MODEL_WRAPPER_REGISTRY.pretrained_models.keys()
232 else:
233 all_model_keys = MODEL_WRAPPER_REGISTRY.registry_dict.keys()
234 all_models = {
235 model_key.model_name + '_' + model_key.use_case + '_' + model_key.framework: model_key
236 for model_key in all_model_keys
237 }
238 models = set()
239 include_filters = (
240 filter_string if isinstance(filter_string, (tuple, list)) else [filter_string]
241 )
242 matched_sets = []
243 for keyword in include_filters:
244 matched = set(fnmatch.filter(all_models.keys(), f'*{keyword}*'))
245 matched_sets.append(matched) # append always, even if empty
246
247 if match_all:
248 # If ANY matched set is empty → intersection is empty
249 if any(len(s) == 0 for s in matched_sets):
250 models = set()
251 else:
252 models = set.intersection(*matched_sets)
253 else:
254 # match_any behavior (if you need it)
255 models = set().union(*matched_sets)
256
257 found_model_keys = [all_models[model] for model in sorted(models)]
258
259 if not print_table:
260 return found_model_keys
261
262 # Build a table with counts per model name (dataset_name removed)
263 model_counts = collections.Counter(mk.model_name for mk in found_model_keys)
264
265 table = texttable.Texttable()
266 rows = [['Model name', 'Count']]
267 for model_name, count in model_counts.items():
268 rows.append([model_name, count])
269
270 table.add_rows(rows)
271 LOGGER.info(table.draw())
272
273 return found_model_keys
274
275
276def list_models_by_dataset(dataset_name, with_checkpoint=False):
277 return [
278 model_key.model_name
279 for model_key in list_models(dataset_name, print_table=False, with_checkpoint=with_checkpoint)
280 if model_key.dataset_name == dataset_name
281 ]
282