Views
No views yet
pip install numpy Pillow huggingface_hub ai-edge-litertclassify.py, paste the script below into it, and save the file1#!/usr/bin/env python3
2import argparse, json
3import numpy as np
4from PIL import Image
5from huggingface_hub import hf_hub_download
6from ai_edge_litert.compiled_model import CompiledModel
7
8def preprocess(img: Image.Image) -> np.ndarray:
9 img = img.convert("RGB")
10 w, h = img.size
11 s = 256
12 if w < h:
13 img = img.resize((s, int(round(h * s / w))), Image.BILINEAR)
14 else:
15 img = img.resize((int(round(w * s / h)), s), Image.BILINEAR)
16 left = (img.size[0] - 224) // 2
17 top = (img.size[1] - 224) // 2
18 img = img.crop((left, top, left + 224, top + 224))
19
20 x = np.asarray(img, dtype=np.float32) / 255.0
21 x = (x - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array(
22 [0.229, 0.224, 0.225], dtype=np.float32
23 )
24 return np.expand_dims(x, axis=0)
25
26def main():
27 ap = argparse.ArgumentParser()
28 ap.add_argument("--image", required=True)
29 args = ap.parse_args()
30
31 model_path = hf_hub_download("litert-community/vgg13", "vgg13.tflite")
32 labels_path = hf_hub_download(
33 "huggingface/label-files", "imagenet-1k-id2label.json", repo_type="dataset"
34 )
35 with open(labels_path, "r", encoding="utf-8") as f:
36 id2label = {int(k): v for k, v in json.load(f).items()}
37
38 img = Image.open(args.image)
39 x = preprocess(img)
40
41 model = CompiledModel.from_file(model_path)
42 inp = model.create_input_buffers(0)
43 out = model.create_output_buffers(0)
44
45 inp[0].write(x)
46 model.run_by_index(0, inp, out)
47
48 req = model.get_output_buffer_requirements(0, 0)
49 y = out[0].read(req["buffer_size"] // np.dtype(np.float32).itemsize, np.float32)
50
51 pred = int(np.argmax(y))
52 label = id2label.get(pred, f"class_{pred}")
53
54 print(f"Top-1 class index: {pred}")
55 print(f"Top-1 label: {label}")
56if __name__ == "__main__":
57 main()python classify.py --image cat.jpg1@misc{simonyan2015deepconvolutionalnetworkslargescale,
2 title={Very Deep Convolutional Networks for Large-Scale Image Recognition},
3 author={Karen Simonyan and Andrew Zisserman},
4 year={2015},
5 eprint={1409.1556},
6 archivePrefix={arXiv},
7 primaryClass={cs.CV},
8 url={https://arxiv.org/abs/1409.1556},
9}