import timm

# 경량 분류 모델 비교
lightweight_models = {
    'mobilenetv3_small_100': '2.5M params, 모바일 최적화',
    'mobilenetv3_large_100': '5.5M params, 모바일 고성능',
    'efficientnet_lite0': '4.7M params, 엣지 최적화',
    'efficientnet_b0': '5.3M params, 범용 경량',
}

for model_name, desc in lightweight_models.items():
    model = timm.create_model(model_name, pretrained=True, num_classes=5)
    params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"{model_name}: {params:.1f}M 파라미터 — {desc}")

# Jetson에서 YOLO 모델 배포
from ultralytics import YOLO

# 1. TensorRT 엔진 생성 (Jetson에서 실행)
model = YOLO('yolo11n.pt')
model.export(
    format='engine',
    imgsz=640,
    half=True,      # FP16 (Jetson GPU 지원)
    device=0,
)

# 2. TensorRT 엔진으로 추론
engine_model = YOLO('yolo11n.engine')
results = engine_model('image.jpg')

# 3. 실시간 카메라 추론
import cv2

cap = cv2.VideoCapture(0)  # CSI 카메라 또는 USB 카메라
while cap.isOpened():
    ret, frame = cap.read()
    if not ret:
        break

    results = engine_model(frame, verbose=False)
    annotated = results[0].plot()

    cv2.imshow('Jetson Inference', annotated)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

# Jetson 환경 확인 명령어
# JetPack 버전 확인
cat /etc/nv_tegra_release

# GPU 사용률 모니터링
tegrastats

# CUDA 버전 확인
nvcc --version

import onnxruntime as ort
import numpy as np
import cv2

# Execution Provider 확인
print(ort.get_available_providers())
# CPU: ['CPUExecutionProvider']
# GPU: ['CUDAExecutionProvider', 'CPUExecutionProvider']

# 세션 생성 (디바이스에 맞는 Provider 자동 선택)
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession('model.onnx', providers=providers)

# 입력 정보 확인
input_info = session.get_inputs()[0]
print(f"입력 이름: {input_info.name}")
print(f"입력 형태: {input_info.shape}")
print(f"입력 타입: {input_info.type}")

# 이미지 전처리
def preprocess(image_path, input_size=(224, 224)):
    """추론용 이미지를 전처리합니다."""
    img = cv2.imread(image_path)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, input_size)
    img = img.astype(np.float32) / 255.0

    # ImageNet 정규화
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img = (img - mean) / std

    # (H, W, C) -> (1, C, H, W)
    img = np.transpose(img, (2, 0, 1))
    img = np.expand_dims(img, axis=0).astype(np.float32)
    return img

# 추론 실행
input_data = preprocess('image.jpg')
outputs = session.run(None, {input_info.name: input_data})
predictions = outputs[0]
predicted_class = np.argmax(predictions, axis=1)[0]
confidence = np.max(predictions)
print(f"예측 클래스: {predicted_class}, 신뢰도: {confidence:.4f}")

# PyTorch → ONNX → TFLite 변환 파이프라인
import torch

# 1. PyTorch 모델 로드
model = timm.create_model('mobilenetv3_small_100', pretrained=True, num_classes=5)
model.eval()

# 2. ONNX 변환
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model, dummy_input, 'mobilenet.onnx',
    input_names=['input'], output_names=['output'],
    opset_version=13,
)

# 3. ONNX → TFLite (onnx2tf 패키지 사용)
# pip install onnx2tf
# onnx2tf -i mobilenet.onnx -o tflite_model

# 4. Ultralytics YOLO는 직접 TFLite Export 지원
from ultralytics import YOLO

yolo = YOLO('yolo11n.pt')
yolo.export(format='tflite', imgsz=320)  # 모바일용 작은 해상도

import time
import numpy as np

def benchmark_inference(model_fn, input_data, n_warmup=10, n_runs=100):
    """추론 성능을 벤치마크합니다."""
    # 워밍업
    for _ in range(n_warmup):
        model_fn(input_data)

    # 측정
    latencies = []
    for _ in range(n_runs):
        start = time.perf_counter()
        model_fn(input_data)
        latencies.append((time.perf_counter() - start) * 1000)

    latencies = np.array(latencies)
    print(f"평균: {latencies.mean():.2f} ms")
    print(f"중간값: {np.median(latencies):.2f} ms")
    print(f"95th: {np.percentile(latencies, 95):.2f} ms")
    print(f"FPS: {1000 / latencies.mean():.1f}")

# ONNX Runtime 벤치마크
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
benchmark_inference(
    lambda x: session.run(None, {'input': x}),
    input_data,
)