File size: 2,790 Bytes
2c83b43 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """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()
|