INNER CODE UNIT · Python
verify_tflite
ultralytics/yolo-flutter-app · scripts/export-tflite-models.py:79
def verify_tflite(path: Path, imgsz: int) -> list[tuple[int, ...]]:
"""Verify the fixed input size, run one zero-input inference, and return output shapes."""
# ai_edge_litert ships with ultralytics[export-litert]; TensorFlow is not installed in that environment
from ai_edge_litert.interpreter import Interpreter
interpreter = Interpreter(model_path=str(path))
interpreter.allocate_tensors()
input_detail = interpreter.get_input_details()[0]
shape = input_detail["shape"]
if list(shape).count(imgsz) != 2:
raise ValueError(f"{path.name} input is {shape.tolist()}; expected two {imgsz}-pixel spatial dimensions")
dtype = input_detail["dtype"]
sample = np.zeros(shape, dtype=dtype)
if not np.issubdtype(dtype, np.floating):
sample.fill(input_detail.get("quantization", (0.0, 0))[1])
interpreter.set_tensor(input_detail["index"], sample)
interpreter.invoke()
return [tuple(detail["shape"].tolist()) for detail in interpreter.get_output_details()]