Views
No views yet
SqueezeNet1_0_Weights.IMAGENET1K_V1).torchvision repository from which it was converted.(B, C, H, W) layout and adds the required Batch dimension to match the LiteRT (B, H, W, C) (NHWC) runtime expectation.pip install numpy Pillow huggingface_hub ai-edge-litertclassify.py, paste the script below into it, and save the file:1#!/usr/bin/env python3
2import argparse
3import json
4import numpy as np
5from PIL import Image
6from huggingface_hub import hf_hub_download
7from ai_edge_litert.compiled_model import CompiledModel
8
9def preprocess(img: Image.Image) -> np.ndarray:
10 img = img.convert("RGB")
11 w, h = img.size
12
13 # Resize shortest edge to 256
14 s = 256
15 if w < h:
16 img = img.resize((s, int(round(h * s / w))), Image.BILINEAR)
17 else:
18 img = img.resize((int(round(w * s / h)), s), Image.BILINEAR)
19
20 # Central crop to 224x224
21 left = (img.size[0] - 224) // 2
22 top = (img.size[1] - 224) // 2
23 img = img.crop((left, top, left + 224, top + 224))
24
25 # Rescale to [0.0, 1.0] and Normalize
26 x = np.asarray(img, dtype=np.float32) / 255.0
27 x = (x - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array(
28 [0.229, 0.224, 0.225], dtype=np.float32
29 )
30
31 # Expand dimensions to create NHWC 4D tensor: (1, 224, 224, 3)
32 x = np.expand_dims(x, axis=0)
33
34 return x
35
36def main():
37 ap = argparse.ArgumentParser()
38 ap.add_argument("--image", required=True, help="Path to the input image")
39 args = ap.parse_args()
40
41 # Download the TFLite model and labels
42 model_path = hf_hub_download("litert-community/squeezenet1_0", "squeezenet1_0.tflite")
43 labels_path = hf_hub_download(
44 "huggingface/label-files", "imagenet-1k-id2label.json", repo_type="dataset"
45 )
46
47 with open(labels_path, "r", encoding="utf-8") as f:
48 id2label = {int(k): v for k, v in json.load(f).items()}
49
50 img = Image.open(args.image)
51 x = preprocess(img)
52
53 model = CompiledModel.from_file(model_path)
54 inp = model.create_input_buffers(0)
55 out = model.create_output_buffers(0)
56
57 inp[0].write(x)
58 model.run_by_index(0, inp, out)
59
60 req = model.get_output_buffer_requirements(0, 0)
61 y = out[0].read(req["buffer_size"] // np.dtype(np.float32).itemsize, np.float32)
62
63 pred = int(np.argmax(y))
64 label = id2label.get(pred, f"class_{pred}")
65
66 print(f"Top-1 class index: {pred}")
67 print(f"Top-1 label: {label}")
68
69if __name__ == "__main__":
70 main()python classify.py --image cat.jpg1@misc{iandola2016squeezenetalexnetlevelaccuracy50x,
2 title={SqueezeNet: AlexNet-level accuracy with 50x fewer parameters and <0.5MB model size},
3 author={Forrest N. Iandola and Song Han and Matthew W. Moskewicz and Khalid Ashraf and William J. Dally and Kurt Keutzer},
4 year={2016},
5 eprint={1602.07360},
6 archivePrefix={arXiv},
7 primaryClass={cs.CV},
8 url={https://arxiv.org/abs/1602.07360},
9}