administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-05-03 20:17:42 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Fix factory

Browse Source
This commit is contained in:
Nicolas Mowen 2025-05-29 14:59:57 -06:00
parent e3af887927
commit 6e2edac90d
1 changed files with 19 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
frigate/util/classification.py
View File

@ -15,7 +15,8 @@ LEARNING_RATE = 0.001
@staticmethod
def generate_representative_dataset(dataset_dir: str):
def generate_representative_dataset_factory(dataset_dir: str):
def generate_representative_dataset():
image_paths = []
for root, dirs, files in os.walk(dataset_dir):
for file in files:
@ -30,6 +31,8 @@ def generate_representative_dataset(dataset_dir: str):
img_array = img_array[None, ...]
yield [img_array]
return generate_representative_dataset
@staticmethod
def train_classification_model(model_dir: str) -> bool:
@ -92,7 +95,9 @@ def train_classification_model(model_dir: str) -> bool:
# convert model to tflite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.representative_dataset = generate_representative_dataset(dataset_dir)
converter.representative_dataset = generate_representative_dataset_factory(
dataset_dir
)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.uint8
converter.inference_output_type = tf.uint8