42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
import tensorflow as tf
|
||
import numpy as np
|
||
import os
|
||
from PIL import Image
|
||
|
||
dir = os.path.dirname(os.path.realpath(__file__))
|
||
|
||
# --- Pfad zum SavedModel ---
|
||
saved_model_dir = dir + "/yolo_training/NAO_detector9/weights/best_saved_model"
|
||
|
||
# --- Optional: Pfad zu Beispielbildern ---
|
||
image_dir = dir + "/calib_images" # z. B. 50–100 JPGs aus deinem Datensatz (z. B. NAO-Roboter-Bilder)
|
||
input_size = (320, 320) # oder (640, 640) je nach deinem Modellinput
|
||
|
||
# --- Repräsentative Datenfunktion ---
|
||
def representative_data_gen():
|
||
for filename in os.listdir(image_dir):
|
||
if filename.endswith(".jpg") or filename.endswith(".png"):
|
||
img = Image.open(os.path.join(image_dir, filename)).convert("RGB")
|
||
img = img.resize(input_size)
|
||
img = np.array(img, dtype=np.float32) / 255.0 # normalisieren falls nötig
|
||
img = np.expand_dims(img, axis=0)
|
||
yield [img]
|
||
|
||
# --- Konverter konfigurieren ---
|
||
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
|
||
converter.optimizations = [tf.lite.Optimize.DEFAULT]
|
||
converter.representative_dataset = representative_data_gen
|
||
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
|
||
|
||
# Je nach Ziel-Hardware (z. B. NAO): INT8 in/out oder UINT8
|
||
converter.inference_input_type = tf.uint8
|
||
converter.inference_output_type = tf.float32
|
||
|
||
# --- Konvertieren ---
|
||
quant_model = converter.convert()
|
||
|
||
# --- Speichern ---
|
||
with open("best_int8.tflite", "wb") as f:
|
||
f.write(quant_model)
|
||
|
||
print("✅ INT8-Quantisierung abgeschlossen.") |