onnx_inference.py

  1#!/usr/bin/env python3
  2"""
  3ONNX Runtime 기반 Edge AI 추론
  4이미지 분류 및 객체 검출 예제
  5
  6참고: content/ko/IoT_Embedded/09_Edge_AI_ONNX.md
  7"""
  8
  9import numpy as np
 10import time
 11from typing import Optional, Tuple, List, Dict
 12import os
 13
 14# ONNX Runtime 설치 확인
 15try:
 16    import onnxruntime as ort
 17    HAS_ONNX = True
 18except ImportError:
 19    HAS_ONNX = False
 20    print("경고: onnxruntime이 설치되지 않았습니다.")
 21    print("설치: pip install onnxruntime")
 22
 23# OpenCV 설치 확인
 24try:
 25    import cv2
 26    HAS_CV2 = True
 27except ImportError:
 28    HAS_CV2 = False
 29    print("경고: opencv-python이 설치되지 않았습니다.")
 30    print("설치: pip install opencv-python")
 31
 32# PIL 설치 확인
 33try:
 34    from PIL import Image
 35    HAS_PIL = True
 36except ImportError:
 37    HAS_PIL = False
 38    print("경고: Pillow가 설치되지 않았습니다.")
 39    print("설치: pip install Pillow")
 40
 41
 42# === ONNX 모델 래퍼 ===
 43
 44class ONNXModel:
 45    """ONNX 모델 기본 래퍼"""
 46
 47    def __init__(self, model_path: str, providers: Optional[List[str]] = None):
 48        if not HAS_ONNX:
 49            raise ImportError("onnxruntime이 필요합니다")
 50
 51        if not os.path.exists(model_path):
 52            raise FileNotFoundError(f"모델 파일을 찾을 수 없습니다: {model_path}")
 53
 54        if providers is None:
 55            # 사용 가능한 프로바이더 자동 선택
 56            available = ort.get_available_providers()
 57            if 'CUDAExecutionProvider' in available:
 58                providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
 59            else:
 60                providers = ['CPUExecutionProvider']
 61
 62        # 세션 옵션 설정
 63        sess_options = ort.SessionOptions()
 64        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
 65        sess_options.intra_op_num_threads = 4  # CPU 스레드 수
 66
 67        # 세션 생성
 68        self.session = ort.InferenceSession(
 69            model_path,
 70            sess_options=sess_options,
 71            providers=providers
 72        )
 73
 74        # 입출력 정보
 75        self.input_name = self.session.get_inputs()[0].name
 76        self.input_shape = self.session.get_inputs()[0].shape
 77        self.input_type = self.session.get_inputs()[0].type
 78        self.output_name = self.session.get_outputs()[0].name
 79
 80        print(f"모델 로드 완료: {model_path}")
 81        print(f"  프로바이더: {self.session.get_providers()}")
 82        print(f"  입력: {self.input_name} {self.input_shape}")
 83        print(f"  출력: {self.output_name}")
 84
 85    def get_input_shape(self) -> list:
 86        """입력 형태 반환"""
 87        return self.input_shape
 88
 89    def predict(self, input_data: np.ndarray) -> np.ndarray:
 90        """추론 수행"""
 91        outputs = self.session.run(
 92            [self.output_name],
 93            {self.input_name: input_data}
 94        )
 95        return outputs[0]
 96
 97    def benchmark(self, num_iterations: int = 100) -> Dict[str, float]:
 98        """성능 벤치마크"""
 99        # 더미 입력 생성
100        dummy_shape = [1 if x == 'batch' or x == 'N' or x is None else x
101                      for x in self.input_shape]
102        dummy_input = np.random.randn(*dummy_shape).astype(np.float32)
103
104        # 워밍업
105        for _ in range(10):
106            self.predict(dummy_input)
107
108        # 측정
109        times = []
110        for _ in range(num_iterations):
111            start = time.perf_counter()
112            self.predict(dummy_input)
113            elapsed = (time.perf_counter() - start) * 1000  # ms
114            times.append(elapsed)
115
116        times = np.array(times)
117
118        results = {
119            "mean_ms": float(np.mean(times)),
120            "std_ms": float(np.std(times)),
121            "min_ms": float(np.min(times)),
122            "max_ms": float(np.max(times)),
123            "fps": 1000.0 / np.mean(times)
124        }
125
126        return results
127
128
129# === 이미지 분류 모델 ===
130
131class ImageClassifier(ONNXModel):
132    """ONNX 이미지 분류 모델"""
133
134    # ImageNet 클래스 (상위 10개만 예시)
135    IMAGENET_CLASSES = [
136        'tench', 'goldfish', 'great_white_shark', 'tiger_shark',
137        'hammerhead', 'electric_ray', 'stingray', 'cock', 'hen', 'ostrich'
138        # ... 실제로는 1000개 클래스
139    ]
140
141    def __init__(self, model_path: str, labels_path: Optional[str] = None):
142        super().__init__(model_path)
143
144        # 레이블 로드 (있는 경우)
145        if labels_path and os.path.exists(labels_path):
146            with open(labels_path, 'r') as f:
147                self.labels = [line.strip() for line in f]
148        else:
149            self.labels = self.IMAGENET_CLASSES
150
151        # 입력 크기 추출
152        self.input_height = self.input_shape[2] if len(self.input_shape) > 2 else 224
153        self.input_width = self.input_shape[3] if len(self.input_shape) > 3 else 224
154
155    def preprocess_image(self, image_path: str) -> np.ndarray:
156        """이미지 전처리 (PIL 사용)"""
157        if not HAS_PIL:
158            raise ImportError("Pillow가 필요합니다")
159
160        # 이미지 로드
161        image = Image.open(image_path).convert('RGB')
162
163        # 리사이즈
164        image = image.resize((self.input_width, self.input_height))
165
166        # NumPy 배열로 변환
167        img_array = np.array(image).astype(np.float32)
168
169        # 정규화 (ImageNet 표준)
170        mean = np.array([0.485, 0.456, 0.406]) * 255
171        std = np.array([0.229, 0.224, 0.225]) * 255
172        img_array = (img_array - mean) / std
173
174        # HWC to CHW
175        img_array = img_array.transpose(2, 0, 1)
176
177        # 배치 차원 추가
178        img_array = np.expand_dims(img_array, axis=0)
179
180        return img_array
181
182    def classify(self, image_path: str, top_k: int = 5) -> List[Tuple[str, float]]:
183        """이미지 분류"""
184        # 전처리
185        input_data = self.preprocess_image(image_path)
186
187        # 추론
188        start = time.perf_counter()
189        output = self.predict(input_data)
190        inference_time = (time.perf_counter() - start) * 1000
191
192        # Softmax
193        probs = self._softmax(output[0])
194
195        # Top-K 결과
196        top_indices = np.argsort(probs)[-top_k:][::-1]
197
198        results = []
199        for idx in top_indices:
200            label = self.labels[idx] if idx < len(self.labels) else f"class_{idx}"
201            results.append((label, float(probs[idx])))
202
203        print(f"추론 시간: {inference_time:.2f}ms")
204
205        return results
206
207    @staticmethod
208    def _softmax(x: np.ndarray) -> np.ndarray:
209        """Softmax 함수"""
210        exp_x = np.exp(x - np.max(x))
211        return exp_x / exp_x.sum()
212
213
214# === 객체 검출 모델 (YOLO) ===
215
216class YOLODetector:
217    """YOLO ONNX 객체 검출기"""
218
219    # COCO 데이터셋 80개 클래스
220    COCO_CLASSES = [
221        'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train',
222        'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign',
223        'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep',
224        'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella',
225        'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
226        'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard',
227        'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork',
228        'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
229        'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair',
230        'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv',
231        'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave',
232        'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase',
233        'scissors', 'teddy bear', 'hair drier', 'toothbrush'
234    ]
235
236    def __init__(self, model_path: str, conf_threshold: float = 0.5,
237                 iou_threshold: float = 0.45):
238        if not HAS_ONNX:
239            raise ImportError("onnxruntime이 필요합니다")
240
241        if not os.path.exists(model_path):
242            # 모델이 없는 경우 시뮬레이션 모드
243            print(f"경고: 모델 파일을 찾을 수 없습니다: {model_path}")
244            print("시뮬레이션 모드로 실행합니다.")
245            self.simulation_mode = True
246            self.input_height = 640
247            self.input_width = 640
248            return
249
250        self.simulation_mode = False
251
252        # ONNX 세션 생성
253        sess_options = ort.SessionOptions()
254        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
255
256        self.session = ort.InferenceSession(
257            model_path,
258            sess_options=sess_options,
259            providers=['CPUExecutionProvider']
260        )
261
262        self.conf_threshold = conf_threshold
263        self.iou_threshold = iou_threshold
264
265        # 입력 정보
266        input_info = self.session.get_inputs()[0]
267        self.input_name = input_info.name
268        self.input_shape = input_info.shape
269        self.input_height = self.input_shape[2]
270        self.input_width = self.input_shape[3]
271
272        print(f"YOLO 모델 로드 완료")
273        print(f"  입력 크기: {self.input_width}x{self.input_height}")
274
275    def preprocess(self, image: np.ndarray) -> Tuple[np.ndarray, Tuple[float, float]]:
276        """이미지 전처리"""
277        orig_height, orig_width = image.shape[:2]
278
279        # 리사이즈
280        resized = cv2.resize(image, (self.input_width, self.input_height))
281
282        # BGR to RGB, HWC to CHW
283        input_data = resized[:, :, ::-1].transpose(2, 0, 1)
284
285        # 정규화 (0-1)
286        input_data = input_data.astype(np.float32) / 255.0
287
288        # 배치 차원 추가
289        input_data = np.expand_dims(input_data, axis=0)
290
291        # 스케일 비율 저장
292        scale = (orig_width / self.input_width, orig_height / self.input_height)
293
294        return input_data, scale
295
296    def detect(self, image: np.ndarray) -> List[Dict]:
297        """객체 검출"""
298        if self.simulation_mode:
299            # 시뮬레이션: 랜덤 검출 결과 반환
300            print("시뮬레이션 모드: 랜덤 검출 결과를 생성합니다.")
301            return self._simulate_detection(image)
302
303        if not HAS_CV2:
304            raise ImportError("opencv-python이 필요합니다")
305
306        # 전처리
307        input_data, scale = self.preprocess(image)
308
309        # 추론
310        start = time.perf_counter()
311        outputs = self.session.run(None, {self.input_name: input_data})
312        inference_time = (time.perf_counter() - start) * 1000
313
314        # 후처리
315        detections = self.postprocess(outputs[0], scale)
316
317        print(f"추론 시간: {inference_time:.2f}ms")
318        print(f"검출된 객체: {len(detections)}개")
319
320        return detections
321
322    def postprocess(self, output: np.ndarray, scale: Tuple[float, float]) -> List[Dict]:
323        """출력 후처리"""
324        if not HAS_CV2:
325            return []
326
327        predictions = output[0]
328
329        boxes = []
330        scores = []
331        class_ids = []
332
333        for pred in predictions:
334            confidence = pred[4]
335
336            if confidence > self.conf_threshold:
337                class_probs = pred[5:]
338                class_id = np.argmax(class_probs)
339                class_score = class_probs[class_id]
340
341                if class_score > self.conf_threshold:
342                    # 박스 좌표 (center_x, center_y, width, height)
343                    cx, cy, w, h = pred[:4]
344
345                    # 원본 스케일로 변환
346                    x1 = int((cx - w / 2) * scale[0])
347                    y1 = int((cy - h / 2) * scale[1])
348                    x2 = int((cx + w / 2) * scale[0])
349                    y2 = int((cy + h / 2) * scale[1])
350
351                    boxes.append([x1, y1, x2, y2])
352                    scores.append(float(confidence * class_score))
353                    class_ids.append(int(class_id))
354
355        # NMS (Non-Maximum Suppression)
356        if boxes:
357            indices = cv2.dnn.NMSBoxes(
358                boxes, scores, self.conf_threshold, self.iou_threshold
359            )
360
361            results = []
362            for i in indices:
363                idx = i[0] if isinstance(i, (list, np.ndarray)) else i
364                results.append({
365                    'box': boxes[idx],
366                    'score': scores[idx],
367                    'class_id': class_ids[idx],
368                    'class_name': self.COCO_CLASSES[class_ids[idx]]
369                })
370
371            return results
372
373        return []
374
375    def _simulate_detection(self, image: np.ndarray) -> List[Dict]:
376        """시뮬레이션: 랜덤 검출 결과"""
377        height, width = image.shape[:2]
378
379        num_detections = np.random.randint(1, 5)
380        detections = []
381
382        for _ in range(num_detections):
383            x1 = np.random.randint(0, width // 2)
384            y1 = np.random.randint(0, height // 2)
385            x2 = np.random.randint(x1 + 50, width)
386            y2 = np.random.randint(y1 + 50, height)
387
388            class_id = np.random.randint(0, len(self.COCO_CLASSES))
389
390            detections.append({
391                'box': [x1, y1, x2, y2],
392                'score': np.random.uniform(0.5, 0.95),
393                'class_id': class_id,
394                'class_name': self.COCO_CLASSES[class_id]
395            })
396
397        return detections
398
399    def draw_detections(self, image: np.ndarray, detections: List[Dict]) -> np.ndarray:
400        """검출 결과 시각화"""
401        if not HAS_CV2:
402            print("경고: opencv-python이 없어 시각화를 건너뜁니다.")
403            return image
404
405        result = image.copy()
406
407        for det in detections:
408            x1, y1, x2, y2 = det['box']
409            label = f"{det['class_name']}: {det['score']:.2f}"
410
411            # 박스 그리기
412            cv2.rectangle(result, (x1, y1), (x2, y2), (0, 255, 0), 2)
413
414            # 라벨 배경
415            (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
416            cv2.rectangle(result, (x1, y1 - 20), (x1 + w, y1), (0, 255, 0), -1)
417
418            # 라벨 텍스트
419            cv2.putText(result, label, (x1, y1 - 5),
420                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
421
422        return result
423
424
425# === 사용 예제 ===
426
427def example_basic_inference():
428    """기본 ONNX 추론 예제"""
429    print("\n=== 기본 ONNX 추론 예제 ===")
430
431    if not HAS_ONNX:
432        print("onnxruntime이 설치되지 않아 예제를 실행할 수 없습니다.")
433        return
434
435    # 시뮬레이션: 더미 모델 생성
436    print("시뮬레이션 모드: 더미 입력으로 테스트합니다.")
437
438    # 더미 데이터
439    batch_size = 1
440    channels = 3
441    height = 224
442    width = 224
443
444    dummy_input = np.random.randn(batch_size, channels, height, width).astype(np.float32)
445
446    print(f"입력 형태: {dummy_input.shape}")
447    print(f"입력 데이터 범위: [{dummy_input.min():.2f}, {dummy_input.max():.2f}]")
448
449
450def example_image_classification():
451    """이미지 분류 예제"""
452    print("\n=== 이미지 분류 예제 ===")
453
454    if not HAS_ONNX:
455        print("onnxruntime이 설치되지 않아 예제를 실행할 수 없습니다.")
456        return
457
458    # 모델 경로 (예시)
459    model_path = "resnet18.onnx"
460
461    if not os.path.exists(model_path):
462        print(f"모델 파일을 찾을 수 없습니다: {model_path}")
463        print("PyTorch에서 변환 예시:")
464        print("  import torch")
465        print("  model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=True)")
466        print("  dummy_input = torch.randn(1, 3, 224, 224)")
467        print("  torch.onnx.export(model, dummy_input, 'resnet18.onnx')")
468        return
469
470    # 분류기 생성
471    classifier = ImageClassifier(model_path)
472
473    # 벤치마크
474    print("\n성능 벤치마크:")
475    results = classifier.benchmark(num_iterations=50)
476    print(f"  평균: {results['mean_ms']:.2f}ms")
477    print(f"  FPS: {results['fps']:.1f}")
478
479
480def example_object_detection():
481    """객체 검출 예제"""
482    print("\n=== 객체 검출 예제 (시뮬레이션) ===")
483
484    # 시뮬레이션 모드로 실행
485    detector = YOLODetector("yolov5s.onnx")  # 파일이 없어도 시뮬레이션 가능
486
487    # 더미 이미지 생성
488    dummy_image = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
489
490    # 검출
491    detections = detector.detect(dummy_image)
492
493    # 결과 출력
494    print("\n검출 결과:")
495    for i, det in enumerate(detections):
496        print(f"  {i+1}. {det['class_name']}: {det['score']:.2f}")
497        print(f"     박스: {det['box']}")
498
499    # 시각화 (OpenCV가 있는 경우)
500    if HAS_CV2:
501        result_image = detector.draw_detections(dummy_image, detections)
502        print("\n결과 이미지 생성 완료")
503
504
505def example_performance_comparison():
506    """성능 비교 예제"""
507    print("\n=== 성능 비교 예제 ===")
508
509    if not HAS_ONNX:
510        print("onnxruntime이 설치되지 않아 예제를 실행할 수 없습니다.")
511        return
512
513    print("배치 크기별 성능 비교 (시뮬레이션)")
514
515    input_shape = (1, 3, 224, 224)
516
517    for batch_size in [1, 4, 8, 16]:
518        data = np.random.randn(batch_size, *input_shape[1:]).astype(np.float32)
519
520        start = time.perf_counter()
521        # 시뮬레이션: 간단한 연산
522        _ = np.mean(data, axis=(2, 3))
523        elapsed = time.perf_counter() - start
524
525        throughput = batch_size / elapsed
526        print(f"배치 크기 {batch_size:2d}: {throughput:.1f} samples/sec")
527
528
529# === 메인 실행 ===
530
531if __name__ == "__main__":
532    print("=" * 60)
533    print("ONNX Runtime Edge AI 추론 예제")
534    print("=" * 60)
535
536    # ONNX Runtime 설치 확인
537    if HAS_ONNX:
538        print(f"\nONNX Runtime 버전: {ort.__version__}")
539        print(f"사용 가능한 프로바이더: {ort.get_available_providers()}")
540    else:
541        print("\n경고: ONNX Runtime이 설치되지 않았습니다.")
542        print("설치: pip install onnxruntime")
543
544    # 예제 실행
545    example_basic_inference()
546    example_image_classification()
547    example_object_detection()
548    example_performance_comparison()
549
550    print("\n" + "=" * 60)
551    print("모든 예제 완료")
552    print("=" * 60)