Full catalogue sync: example.py

example.py

@@ -14,6 +14,7 @@ import argparse
 import json
 import math
 from pathlib import Path
+from torchvision.models import GoogLeNet_Weights
 
 import numpy as np
 from PIL import Image
@@ -22,10 +23,9 @@ from PIL import Image
 # Paths
 # ---------------------------------------------------------------------------
 SCRIPT_DIR = Path(__file__).parent.resolve()
-_model_in_parent = SCRIPT_DIR.parent / "deit-tiny-int8-litert.tflite"
-MODEL_PATH = _model_in_parent if _model_in_parent.exists() else SCRIPT_DIR / "deit-tiny-int8-litert.tflite"
+_model_in_parent = SCRIPT_DIR.parent / "facebook__deit-tiny-patch16-224_litert_optimized.tflite"
+MODEL_PATH = _model_in_parent if _model_in_parent.exists() else SCRIPT_DIR / "facebook__deit-tiny-patch16-224_litert_optimized.tflite"
 IMAGE_PATH = SCRIPT_DIR / "sample_input.jpg"
-CLASSES_PATH = SCRIPT_DIR / "imagenet_classes.json"
 
 # ---------------------------------------------------------------------------
 # Constants
@@ -97,10 +97,11 @@ def postprocess(logits: np.ndarray, top_k: int = TOP_K) -> list[dict]:
     exp = np.exp(logits - logits.max())
     probs = exp / exp.sum()
     top_indices = np.argsort(probs)[::-1][:top_k]
-    with open(CLASSES_PATH) as f:
-        classes = json.load(f)
+    # ImageNet class labels (1000 classes)
+    WEIGHTS = GoogLeNet_Weights.IMAGENET1K_V1
+    classes = WEIGHTS.meta["categories"]
     return [
-        {"rank": i + 1, "class": classes[str(idx)], "probability": float(probs[idx])}
+        {"rank": i + 1, "class": classes[idx], "probability": float(probs[idx])}
         for i, idx in enumerate(top_indices)
     ]