| """Minimal inference example for google/vit-base-patch16-224 (ViT-Base) using ExecuTorch. |
| |
| Loads a quantized INT8 .pte model and runs image classification on a single |
| image, printing the top-5 ImageNet-1k predictions. |
| |
| Preprocessing: resize the shorter edge to 232px (bilinear), center-crop to |
| 224x224, convert to RGB, scale pixel values to [0, 1], then normalize with |
| mean=[0.5, 0.5, 0.5] and std=[0.5, 0.5, 0.5] (the stats reported by the |
| checkpoint's image processor). This matches the resize-232/center-crop-224 |
| recipe used to calibrate and evaluate this model in the Arm optimization |
| pipeline. |
| |
| Requires: executorch, torch, torchvision, pillow, json |
| """ |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from executorch.runtime import Runtime |
| from PIL import Image, ImageDraw, ImageFont |
| from torchvision import transforms |
|
|
| |
| MODEL_PATH = "google__vit-base-patch16-224_raspberry_executorch_optimized.pte" |
| IMAGE_PATH = "sample_input.jpg" |
| LABELS_PATH = "imagenet_classes.json" |
| RESIZE_SIZE = 232 |
| CROP_SIZE = 224 |
| MEAN = (0.5, 0.5, 0.5) |
| STD = (0.5, 0.5, 0.5) |
| TOP_K = 5 |
|
|
| _TRANSFORM = transforms.Compose( |
| [ |
| transforms.Resize(RESIZE_SIZE), |
| transforms.CenterCrop(CROP_SIZE), |
| transforms.ToTensor(), |
| ] |
| ) |
|
|
|
|
| def load_model(model_path: str): |
| """Load an ExecuTorch .pte model and return its forward method. |
| |
| The method is loaded once and reused across calls. |
| """ |
| runtime = Runtime.get() |
| program = runtime.load_program(model_path) |
| return program.load_method("forward") |
|
|
|
|
| def load_labels(labels_path: str) -> dict[str, str]: |
| """Load the ImageNet-1k index-to-label mapping.""" |
| with open(labels_path) as f: |
| return json.load(f) |
|
|
|
|
| def preprocess(image_path: str) -> torch.Tensor: |
| """Load an image and prepare it as a normalized [1, 3, 224, 224] tensor.""" |
| image = Image.open(image_path).convert("RGB") |
| tensor = _TRANSFORM(image) |
|
|
| mean = torch.tensor(MEAN, dtype=torch.float32).view(3, 1, 1) |
| std = torch.tensor(STD, dtype=torch.float32).view(3, 1, 1) |
| tensor = (tensor - mean) / std |
|
|
| return tensor.unsqueeze(0).contiguous() |
|
|
|
|
| def run_inference(method, input_tensor: torch.Tensor) -> torch.Tensor: |
| """Run a forward pass and return the raw logits tensor [1, num_classes].""" |
| outputs = method.execute([input_tensor]) |
| return outputs[0] |
|
|
|
|
| def postprocess(logits: torch.Tensor, labels: dict[str, str]) -> list[dict]: |
| """Apply softmax and return the top-k class predictions.""" |
| if logits.dim() == 1: |
| logits = logits.unsqueeze(0) |
|
|
| probabilities = torch.softmax(logits, dim=-1) |
| k = min(TOP_K, probabilities.shape[-1]) |
| top_scores, top_indices = probabilities.topk(k, dim=-1) |
|
|
| results = [] |
| for score, idx in zip( |
| top_scores[0].tolist(), top_indices[0].tolist(), strict=False |
| ): |
| results.append( |
| { |
| "class_index": idx, |
| "label": labels.get(str(idx), f"class_{idx}"), |
| "probability": round(score, 6), |
| } |
| ) |
| return results |
|
|
|
|
| def save_results(results: list[dict], output_path: Path) -> None: |
| """Save the top-k predictions as JSON next to the script.""" |
| with open(output_path, "w") as f: |
| json.dump(results, f, indent=2) |
| print(f"Predictions saved to: {output_path}") |
|
|
|
|
| def _load_fonts() -> tuple: |
| """Load DejaVuSans fonts, falling back to the built-in bitmap font.""" |
| font_dir = Path("/usr/share/fonts/truetype/dejavu") |
| try: |
| regular = ImageFont.truetype(str(font_dir / "DejaVuSans.ttf"), 14) |
| bold = ImageFont.truetype(str(font_dir / "DejaVuSans-Bold.ttf"), 15) |
| except OSError: |
| regular = ImageFont.load_default() |
| bold = regular |
| return regular, bold |
|
|
|
|
| def save_output_image(image_path: str, results: list[dict], output_path: Path) -> None: |
| """Render the input image alongside its top-k predictions and save as sample_output.jpg. |
| |
| Layout (700 x 320 px): |
| left panel: input image (280x280) |
| right panel: top-5 prediction bars with class names and probabilities |
| """ |
| MARGIN = 20 |
| IMG_SZ = 280 |
| PANEL_W = 360 |
| CANVAS_W = MARGIN + IMG_SZ + MARGIN + PANEL_W + MARGIN |
| CANVAS_H = MARGIN + IMG_SZ + MARGIN |
| TITLE_H = 28 |
| ROW_H = (IMG_SZ - TITLE_H - 8) // TOP_K |
|
|
| BG_COLOR = "#f5f5f5" |
| PANEL_BG = "#ffffff" |
| TOP1_BG = "#e3f2fd" |
| BAR_TOP1 = "#1565c0" |
| BAR_REST = "#90caf9" |
| TEXT_DARK = "#212121" |
| TEXT_GRAY = "#616161" |
| BORDER = "#bdbdbd" |
|
|
| font, font_bold = _load_fonts() |
|
|
| canvas = Image.new("RGB", (CANVAS_W, CANVAS_H), BG_COLOR) |
| draw = ImageDraw.Draw(canvas) |
|
|
| |
| orig = Image.open(image_path).convert("RGB").resize((IMG_SZ, IMG_SZ), Image.LANCZOS) |
| canvas.paste(orig, (MARGIN, MARGIN)) |
| draw.rectangle( |
| [MARGIN, MARGIN, MARGIN + IMG_SZ - 1, MARGIN + IMG_SZ - 1], |
| outline=BORDER, |
| width=1, |
| ) |
|
|
| |
| px = MARGIN + IMG_SZ + MARGIN |
| py = MARGIN |
| draw.rectangle( |
| [px, py, px + PANEL_W, py + IMG_SZ], fill=PANEL_BG, outline=BORDER, width=1 |
| ) |
|
|
| draw.text((px + 10, py + 6), "Top-5 Predictions", font=font_bold, fill=TEXT_DARK) |
| draw.line( |
| [px + 1, py + TITLE_H, px + PANEL_W - 1, py + TITLE_H], fill=BORDER, width=1 |
| ) |
|
|
| max_prob = results[0]["probability"] if results else 1.0 |
| BAR_MAX_W = PANEL_W - 24 |
|
|
| for rank, pred in enumerate(results): |
| ry = py + TITLE_H + 8 + rank * ROW_H |
|
|
| if rank == 0: |
| draw.rectangle([px + 2, ry, px + PANEL_W - 2, ry + ROW_H - 2], fill=TOP1_BG) |
|
|
| badge_text = f"#{rank + 1}" |
| draw.text( |
| (px + 8, ry + 4), |
| badge_text, |
| font=font_bold, |
| fill=BAR_TOP1 if rank == 0 else TEXT_GRAY, |
| ) |
|
|
| label = pred["label"] |
| max_chars = 32 |
| if len(label) > max_chars: |
| label = label[: max_chars - 1] + "…" |
| draw.text((px + 36, ry + 4), label, font=font, fill=TEXT_DARK) |
|
|
| bar_w = max(4, int((pred["probability"] / max_prob) * BAR_MAX_W)) |
| bar_y = ry + ROW_H - 14 |
| bar_color = BAR_TOP1 if rank == 0 else BAR_REST |
| draw.rectangle([px + 8, bar_y, px + 8 + bar_w, bar_y + 8], fill=bar_color) |
|
|
| pct_text = f"{pred['probability'] * 100:.2f}%" |
| draw.text((px + PANEL_W - 56, ry + 4), pct_text, font=font, fill=TEXT_GRAY) |
|
|
| canvas.save(str(output_path), quality=95) |
| print(f"Output image saved to: {output_path}") |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--output", type=Path, help="Write predictions as JSON") |
| parser.add_argument("--output-image", type=Path, help="Write the rendered result image") |
| args = parser.parse_args() |
|
|
| script_dir = Path(__file__).resolve().parent |
| model_path = str(script_dir / MODEL_PATH) |
| image_path = str(script_dir / IMAGE_PATH) |
| labels_path = str(script_dir / LABELS_PATH) |
|
|
| method = load_model(model_path) |
| labels = load_labels(labels_path) |
|
|
| input_tensor = preprocess(image_path) |
| logits = run_inference(method, input_tensor) |
| results = postprocess(logits, labels) |
|
|
| print(f"Top-{TOP_K} predictions for {IMAGE_PATH}:") |
| for rank, pred in enumerate(results, start=1): |
| print(f" {rank}. {pred['label']} ({pred['probability']:.4f})") |
|
|
| if args.output: |
| save_results(results, args.output) |
| if args.output_image: |
| save_output_image(image_path, results, args.output_image) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
