#边缘计算：树莓派、手机端与边缘AI部署详解

import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
import time

class RaspberryPiMobileNet:
    def __init__(self):
        # 加载预训练的 MobileNetV2，切换到 CPU 并设为评估模式
        self.device = torch.device("cpu")
        self.model = models.mobilenet_v2(pretrained=True).to(self.device).eval()
        
        # 图像预处理：与训练 ImageNet 时的操作保持一致
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
        
        # 读取 ImageNet 的类别标签
        with open("imagenet_labels.txt", "r") as f:
            self.labels = [line.strip() for line in f.readlines()]

    def predict(self, image_path: str) -> dict:
        # 读取图像并预处理
        img = Image.open(image_path).convert("RGB")
        input_tensor = self.transform(img).unsqueeze(0).to(self.device)
        
        # 计时推理（不计算梯度，省内存）
        start = time.time()
        with torch.no_grad():
            outputs = self.model(input_tensor)
        inference_ms = (time.time() - start) * 1000
        
        # 取出概率最高的3个类别
        probs = torch.nn.functional.softmax(outputs[0], dim=0)
        top3_probs, top3_indices = torch.topk(probs, 3)
        
        return {
            "top3": [(self.labels[i], round(p.item(), 3))
                     for i, p in zip(top3_indices, top3_probs)],
            "inference_ms": round(inference_ms, 2)
        }

if __name__ == "__main__":
    detector = RaspberryPiMobileNet()
    result = detector.predict("test_dog.jpg")
    print(f"预测 Top3: {result['top3']}")
    print(f"推理耗时: {result['inference_ms']}ms")

import torch
import onnx
import tensorflow as tf
import tf2onnx

def pytorch_to_tflite(pytorch_model_path: str, output_tflite_path: str):
    """
    将 PyTorch 模型转为 TensorFlow Lite 格式（带动态量化）
    """
    # 1. 导出 ONNX
    dummy_input = torch.randn(1, 3, 224, 224)
    pytorch_model = models.mobilenet_v2(pretrained=True).eval()
    torch.onnx.export(
        pytorch_model, dummy_input, "temp.onnx",
        export_params=True, opset_version=12,
        do_constant_folding=True
    )
    
    # 2. 校验 ONNX 模型
    onnx_model = onnx.load("temp.onnx")
    onnx.checker.check_model(onnx_model)
    
    # 3. 转换为 TFLite 并启用量化优化
    converter = tf.lite.TFLiteConverter.from_onnx("temp.onnx")
    converter.optimizations = [tf.lite.Optimize.DEFAULT]  # 动态量化
    tflite_model = converter.convert()
    
    # 4. 保存
    with open(output_tflite_path, "wb") as f:
        f.write(tflite_model)
    
    print(f"转换完成，TFLite 模型已保存至：{output_tflite_path}")