Sync model repo (text/metadata)

067cfd6 verified about 2 months ago

7 kB

"""Minimal inference example for microsoft/resnet-50 INT8 using ExecuTorch.
Loads a quantized .pte model and runs image classification on a single image.
Outputs top-5 predictions with class labels and probabilities.
"""

import json
from pathlib import Path

import torch
from executorch.runtime import Runtime
from PIL import Image, ImageDraw, ImageFont
from torchvision import transforms


# ── Configuration ──────────────────────────────────────────────────────────────
MODEL_PATH = "resnet-50_raspberry_executorch_optimized.pte"
IMAGE_PATH = "sample_input.jpg"
INPUT_SIZE = (224, 224)
TOP_K = 5

# ImageNet normalization constants
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

SCRIPT_DIR = Path(__file__).resolve().parent

# ImageNet class labels (1000 classes)
with (SCRIPT_DIR / "imagenet_classes.json").open(encoding="utf-8") as file:
    IMAGENET_CLASSES = json.load(file)


# ── Model Loading ──────────────────────────────────────────────────────────────
def load_model(pte_path: str):
    """Load ExecuTorch .pte model and return the forward method."""
    runtime = Runtime.get()
    program = runtime.load_program(pte_path)
    return program.load_method("forward")


# ── Preprocessing ──────────────────────────────────────────────────────────────
def preprocess(image_path: str) -> torch.Tensor:
    """Load and preprocess image for model input.
    Resizes to 256, center-crops to 224x224, converts to float32 tensor in
    [0, 1], then applies ImageNet mean/std normalization.
    """
    image = Image.open(image_path).convert("RGB")
    transform = transforms.Compose(
        [
            transforms.Resize(256),
            transforms.CenterCrop(INPUT_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
        ]
    )
    return transform(image).unsqueeze(0)  # [1, 3, 224, 224]


# ── Inference ──────────────────────────────────────────────────────────────────
def run_inference(method, input_tensor: torch.Tensor) -> torch.Tensor:
    """Run forward pass and return raw logits tensor."""
    outputs = method.execute([input_tensor])
    return outputs[0]  # [1, 1000]


# ── Postprocessing ─────────────────────────────────────────────────────────────
def postprocess(raw_output: torch.Tensor, labels: dict[str, str], top_k: int = TOP_K) -> list:
    """Decode raw logits into top-k class predictions with probabilities.
    Applies softmax to convert logits to probabilities, then returns the
    top-k classes sorted by descending probability.
    """
    # Ensure 2D: [1, 1000] → [1000]
    logits = raw_output
    if logits.dim() > 1:
        logits = logits.squeeze(0)

    probabilities = torch.softmax(logits, dim=-1)
    k = min(top_k, logits.shape[-1])
    top_scores, top_indices = probabilities.topk(k)

    # Quantized logits land on a discrete grid, so exact ties are common and
    # topk orders them arbitrarily. Break ties by class id for a stable result.
    ranked = sorted(
        zip(top_scores.tolist(), top_indices.tolist(), strict=False),
        key=lambda item: (-item[0], item[1]),
    )

    results = []
    for score, idx in ranked:
        class_label = labels[str(idx)]
        results.append(
            {"class_index": idx, "class": class_label, "probability": round(score, 6)}
        )
    return results


# ── Save Results ───────────────────────────────────────────────────────────────
def save_results(results: list, script_dir: Path) -> None:
    """Save top-k predictions as predictions.json in the script directory."""
    output_path = script_dir / "predictions.json"
    with open(output_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Predictions saved to: {output_path}")


def save_output_image(image_path: str, top_prediction: dict, script_dir: Path) -> None:
    """Overlay the top prediction label on the original image and save as sample_output.jpg."""
    image = Image.open(image_path).convert("RGB")
    draw = ImageDraw.Draw(image)

    label = f"{top_prediction['class']}: {top_prediction['probability'] * 100:.1f}%"

    font_size = max(16, image.width // 20)
    font: ImageFont.ImageFont | ImageFont.FreeTypeFont
    try:
        font = ImageFont.truetype(
            "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size
        )
    except OSError:
        try:
            font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", font_size)
        except OSError:
            font = ImageFont.load_default()

    bbox = draw.textbbox((0, 0), label, font=font)
    text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
    padding = 8
    x, y = padding, image.height - text_h - padding * 3

    draw.rectangle(
        [x - padding, y - padding, x + text_w + padding, y + text_h + padding],
        fill=(0, 0, 0, 180),
    )
    draw.text((x, y), label, fill=(255, 255, 255), font=font)

    output_path = script_dir / "sample_output.jpg"
    image.save(output_path, quality=95)
    print(f"Output image saved to: {output_path}")


# ── Main ───────────────────────────────────────────────────────────────────────
def main() -> None:
    script_dir = Path(__file__).parent

    # Resolve paths relative to script location
    model_path = script_dir / MODEL_PATH
    image_path = script_dir / IMAGE_PATH

    # Load model
    method = load_model(str(model_path))

    # Preprocess input image
    input_tensor = preprocess(str(image_path))

    # Run inference
    raw_output = run_inference(method, input_tensor)

    # Decode predictions
    results = postprocess(raw_output, IMAGENET_CLASSES)

    # Print top-k predictions
    print(f"\nTop-{TOP_K} predictions for: {image_path.name}")
    print("-" * 50)
    for rank, pred in enumerate(results, 1):
        print(f"  {rank}. {pred['class']:<40} {pred['probability']*100:.2f}%")

    # Save predictions.json
    save_results(results, script_dir)

    # Save labeled output image
    save_output_image(str(image_path), results[0], script_dir)


if __name__ == "__main__":
    main()

	"""Minimal inference example for microsoft/resnet-50 INT8 using ExecuTorch.

	Loads a quantized .pte model and runs image classification on a single image.
	Outputs top-5 predictions with class labels and probabilities.
	"""

	import json
	from pathlib import Path

	import torch
	from executorch.runtime import Runtime
	from PIL import Image, ImageDraw, ImageFont
	from torchvision import transforms


	# ── Configuration ──────────────────────────────────────────────────────────────
	MODEL_PATH = "resnet-50_raspberry_executorch_optimized.pte"
	IMAGE_PATH = "sample_input.jpg"
	INPUT_SIZE = (224, 224)
	TOP_K = 5

	# ImageNet normalization constants
	IMAGENET_MEAN = [0.485, 0.456, 0.406]
	IMAGENET_STD = [0.229, 0.224, 0.225]

	SCRIPT_DIR = Path(__file__).resolve().parent

	# ImageNet class labels (1000 classes)
	with (SCRIPT_DIR / "imagenet_classes.json").open(encoding="utf-8") as file:
	IMAGENET_CLASSES = json.load(file)


	# ── Model Loading ──────────────────────────────────────────────────────────────
	def load_model(pte_path: str):
	"""Load ExecuTorch .pte model and return the forward method."""
	runtime = Runtime.get()
	program = runtime.load_program(pte_path)
	return program.load_method("forward")


	# ── Preprocessing ──────────────────────────────────────────────────────────────
	def preprocess(image_path: str) -> torch.Tensor:
	"""Load and preprocess image for model input.

	Resizes to 256, center-crops to 224x224, converts to float32 tensor in
	[0, 1], then applies ImageNet mean/std normalization.
	"""
	image = Image.open(image_path).convert("RGB")
	transform = transforms.Compose(
	[
	transforms.Resize(256),
	transforms.CenterCrop(INPUT_SIZE),
	transforms.ToTensor(),
	transforms.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
	]
	)
	return transform(image).unsqueeze(0) # [1, 3, 224, 224]


	# ── Inference ──────────────────────────────────────────────────────────────────
	def run_inference(method, input_tensor: torch.Tensor) -> torch.Tensor:
	"""Run forward pass and return raw logits tensor."""
	outputs = method.execute([input_tensor])
	return outputs[0] # [1, 1000]


	# ── Postprocessing ─────────────────────────────────────────────────────────────
	def postprocess(raw_output: torch.Tensor, labels: dict[str, str], top_k: int = TOP_K) -> list:
	"""Decode raw logits into top-k class predictions with probabilities.

	Applies softmax to convert logits to probabilities, then returns the
	top-k classes sorted by descending probability.
	"""
	# Ensure 2D: [1, 1000] → [1000]
	logits = raw_output
	if logits.dim() > 1:
	logits = logits.squeeze(0)

	probabilities = torch.softmax(logits, dim=-1)
	k = min(top_k, logits.shape[-1])
	top_scores, top_indices = probabilities.topk(k)

	# Quantized logits land on a discrete grid, so exact ties are common and
	# topk orders them arbitrarily. Break ties by class id for a stable result.
	ranked = sorted(
	zip(top_scores.tolist(), top_indices.tolist(), strict=False),
	key=lambda item: (-item[0], item[1]),
	)

	results = []
	for score, idx in ranked:
	class_label = labels[str(idx)]
	results.append(
	{"class_index": idx, "class": class_label, "probability": round(score, 6)}
	)
	return results


	# ── Save Results ───────────────────────────────────────────────────────────────
	def save_results(results: list, script_dir: Path) -> None:
	"""Save top-k predictions as predictions.json in the script directory."""
	output_path = script_dir / "predictions.json"
	with open(output_path, "w") as f:
	json.dump(results, f, indent=2)
	print(f"Predictions saved to: {output_path}")


	def save_output_image(image_path: str, top_prediction: dict, script_dir: Path) -> None:
	"""Overlay the top prediction label on the original image and save as sample_output.jpg."""
	image = Image.open(image_path).convert("RGB")
	draw = ImageDraw.Draw(image)

	label = f"{top_prediction['class']}: {top_prediction['probability'] * 100:.1f}%"

	font_size = max(16, image.width // 20)
	font: ImageFont.ImageFont \| ImageFont.FreeTypeFont
	try:
	font = ImageFont.truetype(
	"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", font_size
	)
	except OSError:
	try:
	font = ImageFont.truetype("/System/Library/Fonts/Helvetica.ttc", font_size)
	except OSError:
	font = ImageFont.load_default()

	bbox = draw.textbbox((0, 0), label, font=font)
	text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
	padding = 8
	x, y = padding, image.height - text_h - padding * 3

	draw.rectangle(
	[x - padding, y - padding, x + text_w + padding, y + text_h + padding],
	fill=(0, 0, 0, 180),
	)
	draw.text((x, y), label, fill=(255, 255, 255), font=font)

	output_path = script_dir / "sample_output.jpg"
	image.save(output_path, quality=95)
	print(f"Output image saved to: {output_path}")


	# ── Main ───────────────────────────────────────────────────────────────────────
	def main() -> None:
	script_dir = Path(__file__).parent

	# Resolve paths relative to script location
	model_path = script_dir / MODEL_PATH
	image_path = script_dir / IMAGE_PATH

	# Load model
	method = load_model(str(model_path))

	# Preprocess input image
	input_tensor = preprocess(str(image_path))

	# Run inference
	raw_output = run_inference(method, input_tensor)

	# Decode predictions
	results = postprocess(raw_output, IMAGENET_CLASSES)

	# Print top-k predictions
	print(f"\nTop-{TOP_K} predictions for: {image_path.name}")
	print("-" * 50)
	for rank, pred in enumerate(results, 1):
	print(f" {rank}. {pred['class']:<40} {pred['probability']*100:.2f}%")

	# Save predictions.json
	save_results(results, script_dir)

	# Save labeled output image
	save_output_image(str(image_path), results[0], script_dir)


	if __name__ == "__main__":
	main()