inception_v3

#!/usr/bin/env python3
import argparse, json
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download
from ai_edge_litert.compiled_model import CompiledModel

def preprocess(img: Image.Image) -> np.ndarray:
    img = img.convert("RGB")
    w, h = img.size
    # Inception_v3 expects a resize to 342 prior to the 299 central crop
    s = 342
    if w < h:
        img = img.resize((s, int(round(h * s / w))), Image.BILINEAR)
    else:
        img = img.resize((int(round(w * s / h)), s), Image.BILINEAR)
        
    # Central crop to 299x299
    left = (img.size[0] - 299) // 2
    top = (img.size[1] - 299) // 2
    img = img.crop((left, top, left + 299, top + 299))
    
    # Rescale to [0.0, 1.0] and Normalize
    x = np.asarray(img, dtype=np.float32) / 255.0
    x = (x - np.array([0.485, 0.456, 0.406], dtype=np.float32)) / np.array(
        [0.229, 0.224, 0.225], dtype=np.float32
    )
    # Expand dimensions to create NHWC 4D tensor: (1, 299, 299, 3)
    x = np.expand_dims(x, axis=0)

    return x

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", required=True)
    args = ap.parse_args()

    # Download the TFLite model and labels
    model_path = hf_hub_download("litert-community/inception_v3", "inception_v3.tflite")
    labels_path = hf_hub_download(
        "huggingface/label-files", "imagenet-1k-id2label.json", repo_type="dataset"
    )
    
    with open(labels_path, "r", encoding="utf-8") as f:
        id2label = {int(k): v for k, v in json.load(f).items()}

    img = Image.open(args.image)
    x = preprocess(img)

    model = CompiledModel.from_file(model_path)
    inp = model.create_input_buffers(0)
    out = model.create_output_buffers(0)

    inp[0].write(x)
    model.run_by_index(0, inp, out)

    req = model.get_output_buffer_requirements(0, 0)
    y = out[0].read(req["buffer_size"] // np.dtype(np.float32).itemsize, np.float32)

    pred = int(np.argmax(y))
    label = id2label.get(pred, f"class_{pred}")

    print(f"Top-1 class index: {pred}")
    print(f"Top-1 label: {label}")

if __name__ == "__main__":
    main()

@inproceedings{szegedy2016rethinking,
  title={Rethinking the inception architecture for computer vision},
  author={Szegedy, Christian and Vanhoucke, Vincent and Ioffe, Sergey and Shlens, Jon and Wojna, Zbigniew},
  booktitle={Proceedings of the IEEE conference on computer vision and pattern recognition},
  pages={2818--2826},
  year={2016}
}