| """Classify an image on CPU with a packaged FP32, A16W8 or A8W8 TFLite model.""" |
| from pathlib import Path |
| import argparse |
| import json |
|
|
| import numpy as np |
| from PIL import Image |
| from transformers import AutoImageProcessor |
| from ai_edge_litert.interpreter import Interpreter |
|
|
|
|
| def quantization(detail): |
| params = detail["quantization_parameters"] |
| if len(params["scales"]) != 1 or len(params["zero_points"]) != 1: |
| raise ValueError("This example requires per-tensor input/output quantization") |
| scale = float(params["scales"][0]) |
| if not np.isfinite(scale) or scale <= 0: |
| raise ValueError("Invalid input/output quantization scale") |
| return scale, int(params["zero_points"][0]) |
|
|
|
|
| def main(): |
| directory = Path(__file__).resolve().parent |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--image", type=Path, required=True) |
| parser.add_argument("--model", type=Path, default=directory / "model_fp32.tflite", |
| help="TFLite model path; defaults to the packaged FP32 model") |
| args = parser.parse_args() |
| processor = AutoImageProcessor.from_pretrained(directory, backend="pil", local_files_only=True) |
| with Image.open(args.image) as image: |
| pixels = processor(images=image.convert("RGB"), return_tensors="np")["pixel_values"] |
| runtime = Interpreter(model_path=str(args.model), num_threads=4) |
| runtime.allocate_tensors() |
| input_detail, = runtime.get_input_details() |
| output_detail, = runtime.get_output_details() |
| pixels = np.ascontiguousarray(pixels, dtype=np.float32) |
| if tuple(input_detail["shape"]) != pixels.shape: |
| raise ValueError(f"Processor/model shape mismatch: {pixels.shape} vs {input_detail['shape']}") |
| if input_detail["dtype"] == np.int8: |
| scale, zero = quantization(input_detail) |
| pixels = np.clip(np.rint(pixels / scale + zero), -128, 127).astype(np.int8) |
| elif input_detail["dtype"] != np.float32: |
| raise ValueError("Expected FP32 or INT8 input") |
| runtime.set_tensor(input_detail["index"], pixels) |
| runtime.invoke() |
| logits = runtime.get_tensor(output_detail["index"])[0] |
| if output_detail["dtype"] == np.int8: |
| scale, zero = quantization(output_detail) |
| logits = (logits.astype(np.float32) - zero) * scale |
| elif output_detail["dtype"] != np.float32: |
| raise ValueError("Expected FP32 or INT8 output") |
| if logits.shape != (1000,) or not np.isfinite(logits).all(): |
| raise ValueError("Expected 1000 finite classifier logits") |
| labels = json.loads((directory / "labels.json").read_text()) |
| for index in np.argsort(-logits, kind="stable")[:5]: |
| print(f"{int(index)}\t{labels[str(int(index))]}\tlogit={float(logits[index]):.6f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
