compare model extensions (pt,onnx,bin(openvino))

compare_pt_onnx_openvino.py

import cv2
from ultralytics import YOLO
import time
import os
from collections import deque
from pathlib import Path

# Put these files in same dir for easy testing
# - compare_pt_onnx_openvino.ipynb
# - best.pt
# - best.onnx
# - best_openvino_model.bin
# - best_openvino_model.xml
# - best_openvino_model_int8.bin
# - best_openvino_model_int8.xml
# - ros2bag.mp4


# --- 설정 ---
MODEL_BASE_PATH = "/home/f2j/python/onnx/best"  # 확장자 제외한 경로
VIDEO_PATH = "ros2bag.mp4"
CLASS_NAMES = ['crosswalk', 'g', 'p', 'r', 's', 't']

# 테스트 설정
TEST_DEVICE = 'cpu'  # 'cpu' 또는 'cuda' (GPU 테스트 권장!)
USE_CPU_ONLY = True  # True면 CPU 강제, False면 GPU 사용 허용

# CPU 강제 설정 (USE_CPU_ONLY가 True일 때만)
if USE_CPU_ONLY:
    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
    os.environ["ONNXRUNTIME_DISABLE_CUDA"] = "1"
    TEST_DEVICE = 'cpu'

WARMUP_FRAMES = 30  # 워밍업 프레임 수 (초기 로딩 시간 제외)
TEST_FRAMES = 300   # 측정할 프레임 수 (0이면 전체 비디오)

# 테스트할 모델 타입 선택 (True/False로 선택)
TEST_MODELS = {
    'OpenVINO': True,
    'OpenVINO_INT': True,
    'ONNX': True,
    'PT': True,
}

class FPSCounter:
    """FPS 계산 및 통계를 관리하는 클래스"""
    def __init__(self, window_size=30):
        self.frame_times = deque(maxlen=window_size)
        self.total_time = 0
        self.frame_count = 0
        
    def update(self, elapsed_time):
        """프레임 처리 시간 업데이트"""
        self.frame_times.append(elapsed_time)
        self.total_time += elapsed_time
        self.frame_count += 1
    
    def get_current_fps(self):
        """현재 FPS (최근 window_size 프레임 기준)"""
        if len(self.frame_times) == 0:
            return 0
        return len(self.frame_times) / sum(self.frame_times)
    
    def get_average_fps(self):
        """전체 평균 FPS"""
        if self.total_time == 0:
            return 0
        return self.frame_count / self.total_time
    
    def get_stats(self):
        """상세 통계 반환"""
        if self.frame_count == 0:
            return None
        
        avg_fps = self.get_average_fps()
        avg_time = self.total_time / self.frame_count * 1000  # ms
        
        return {
            'total_frames': self.frame_count,
            'total_time': self.total_time,
            'avg_fps': avg_fps,
            'avg_inference_time_ms': avg_time
        }

def get_model_path(base_path: str, model_type: str):
    """모델 타입에 따른 경로 반환"""
    if model_type == 'ONNX':
        return f"{base_path}.onnx"
    elif model_type == 'PT':
        return f"{base_path}.pt"
    elif model_type == 'OpenVINO_INT':
        # OpenVINO는 폴더 구조 (best_openvino_model/)
        openvino_dir = f"{base_path}_openvino_model_int8"
        if os.path.exists(openvino_dir):
            return openvino_dir
        # 또는 .xml 파일
        return f"{base_path}_openvino_model_int8.xml"
    elif model_type == 'OpenVINO':
        # OpenVINO는 폴더 구조 (best_openvino_model/)
        openvino_dir = f"{base_path}_openvino_model"
        if os.path.exists(openvino_dir):
            return openvino_dir
        # 또는 .xml 파일
        return f"{base_path}_openvino_model.xml"
    return None

def check_model_exists(base_path: str, model_type: str):
    """모델 파일/폴더 존재 여부 확인"""
    model_path = get_model_path(base_path, model_type)
    
    if model_type == 'OpenVINO':
        # OpenVINO는 폴더 또는 .xml 파일 체크
        if os.path.isdir(model_path):
            return True, model_path
        xml_path = f"{base_path}.xml"
        if os.path.exists(xml_path):
            return True, xml_path
        return False, None
    elif model_type == "OpenVINO_INT":
        if os.path.isdir(model_path):
            return True, model_path
        xml_path = f"{base_path}.xml"
        if os.path.exists(xml_path):
            return True, xml_path
        return False, None
    
    return os.path.exists(model_path), model_path

def export_openvino_int_if_needed(pt_path: str, base_path: str):
    """OpenVINO 모델이 없으면 PT에서 변환"""
    exists, openvino_path = check_model_exists(base_path, 'OpenVINO_INT')
    
    if exists:
        print(f"✅ OpenVINO 모델이 이미 존재합니다: {openvino_path}")
        return openvino_path
    
    print(f"⚠️  OpenVINO 모델을 찾을 수 없습니다. PT 모델에서 변환 중...")
    try:
        model = YOLO(pt_path)
        export_path = model.export(format='openvino')
        print(f"✅ OpenVINO 모델 변환 완료: {export_path}")
        return export_path
    except Exception as e:
        print(f"❌ OpenVINO 변환 실패: {e}")
        return None
def export_openvino_if_needed(pt_path: str, base_path: str):
    """OpenVINO 모델이 없으면 PT에서 변환"""
    exists, openvino_path = check_model_exists(base_path, 'OpenVINO')
    
    if exists:
        print(f"✅ OpenVINO 모델이 이미 존재합니다: {openvino_path}")
        return openvino_path
    
    print(f"⚠️  OpenVINO 모델을 찾을 수 없습니다. PT 모델에서 변환 중...")
    try:
        model = YOLO(pt_path)
        export_path = model.export(format='openvino')
        print(f"✅ OpenVINO 모델 변환 완료: {export_path}")
        return export_path
    except Exception as e:
        print(f"❌ OpenVINO 변환 실패: {e}")
        return None

def test_model(model_path: str, model_type: str, video_path: str, 
               warmup_frames: int, test_frames: int, device: str = 'cpu', 
               show_video: bool = False):
    """
    단일 모델에 대한 FPS 테스트 수행
    
    Args:
        model_path: 모델 파일 경로
        model_type: 모델 타입 (ONNX, PT, OpenVINO)
        video_path: 비디오 파일 경로
        warmup_frames: 워밍업 프레임 수
        test_frames: 테스트 프레임 수 (0이면 전체)
        device: 사용할 디바이스 ('cpu' 또는 'cuda')
        show_video: 비디오 표시 여부
    """
    print(f"\n{'='*60}")
    print(f"🔍 [{model_type}] 모델 테스트 시작: {model_path}")
    print(f"🖥️  디바이스: {device.upper()}")
    print(f"{'='*60}")
    
    # 1. 모델 로드
    try:
        model = YOLO(model_path, task="detect")
        print(f"✅ 모델 로드 성공")
    except Exception as e:
        print(f"❌ 모델 로드 실패: {e}")
        return None
    
    # 2. 비디오 열기
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"❌ 비디오 파일 로드 실패: {video_path}")
        return None
    
    fps_counter = FPSCounter()
    frame_idx = 0
    
    print(f"\n⏳ 워밍업 중... ({warmup_frames} 프레임)")
    
    # 3. 프레임별 처리
    while cap.isOpened():
        success, frame = cap.read()
        if not success:
            break
        
        frame_idx += 1
        
        # 워밍업 단계
        if frame_idx <= warmup_frames:
            try:
                _ = model.predict(
                    source=frame,
                    conf=0.25,
                    iou=0.45,
                    verbose=False,
                    device=device
                )
            except Exception as e:
                print(f"❌ 추론 중 오류: {e}")
                break
            
            if frame_idx == warmup_frames:
                print(f"✅ 워밍업 완료. 측정 시작...\n")
            continue
        
        # 테스트 프레임 수 체크
        if test_frames > 0 and (frame_idx - warmup_frames) > test_frames:
            break
        
        # 추론 시간 측정
        start_time = time.time()
        try:
            results = model.predict(
                source=frame,
                conf=0.25,
                iou=0.45,
                verbose=False,
                device=device
            )
        except Exception as e:
            print(f"❌ 추론 중 오류: {e}")
            break
        
        elapsed_time = time.time() - start_time
        fps_counter.update(elapsed_time)
        
        # 비디오 표시
        if show_video:
            annotated_frame = results[0].plot()
            
            # FPS 정보 오버레이
            current_fps = fps_counter.get_current_fps()
            cv2.putText(annotated_frame, f"{model_type} | FPS: {current_fps:.1f}", 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
            cv2.putText(annotated_frame, f"Frame: {frame_idx - warmup_frames}/{test_frames if test_frames > 0 else '?'}", 
                       (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
            
            cv2.imshow(f"YOLOv8 {model_type} Performance Test", annotated_frame)
            
            if cv2.waitKey(1) & 0xFF == ord('q'):
                print("✋ 사용자 요청으로 테스트 중단")
                break
        
        # 진행률 표시 (10% 단위)
        if test_frames > 0:
            progress = (frame_idx - warmup_frames) / test_frames * 100
            if (frame_idx - warmup_frames) % max(1, test_frames // 10) == 0:
                print(f"진행률: {progress:.1f}% | 현재 FPS: {fps_counter.get_current_fps():.2f}")
    
    # 4. 자원 해제
    cap.release()
    if show_video:
        cv2.destroyAllWindows()
    
    # 5. 결과 반환
    stats = fps_counter.get_stats()
    if stats:
        print(f"\n📊 [{model_type}] 테스트 결과:")
        print(f"  - 처리 프레임: {stats['total_frames']}")
        print(f"  - 총 처리 시간: {stats['total_time']:.2f}초")
        print(f"  - 평균 FPS: {stats['avg_fps']:.2f}")
        print(f"  - 평균 추론 시간: {stats['avg_inference_time_ms']:.2f}ms")
    
    return stats

def compare_models(base_path: str, video_path: str, device: str = 'cpu',
                   warmup_frames: int = 30, test_frames: int = 300, 
                   test_models: dict = None):
    """
    ONNX, PT, OpenVINO 모델 성능 비교
    """
    if test_models is None:
        test_models = TEST_MODELS
    
    print("🚀 YOLO 모델 성능 비교 테스트")
    print(f"📹 비디오: {video_path}")
    print(f"🖥️  디바이스: {device.upper()}")
    print(f"🔥 워밍업: {warmup_frames} 프레임")
    print(f"📊 테스트: {test_frames if test_frames > 0 else '전체'} 프레임")
    print(f"🎯 테스트 모델: {', '.join([k for k, v in test_models.items() if v])}")
    
    results = {}
    
    # PT 모델 경로 (OpenVINO 변환에 필요)
    pt_path = get_model_path(base_path, 'PT')
    
    # 각 모델 타입별 테스트
    for model_type, should_test in test_models.items():
        if not should_test:
            continue
        
        # 모델 경로 확인
        if model_type == 'OpenVINO':
            # OpenVINO는 필요시 변환
            if not os.path.exists(pt_path):
                print(f"\n⚠️  [{model_type}] PT 모델이 없어 변환할 수 없습니다. 건너뜁니다.")
                continue
            model_path = export_openvino_if_needed(pt_path, base_path)
            if model_path is None:
                continue
        elif model_type == 'OpenVINO_INT':
            if not os.path.exists(pt_path):
                print(f"\n⚠️  [{model_type}] PT 모델이 없어 변환할 수 없습니다. 건너뜁니다.")
                continue
            model_path = export_openvino_int_if_needed(pt_path, base_path)
            if model_path is None:
                continue

        else:
            exists, model_path = check_model_exists(base_path, model_type)
            if not exists:
                print(f"\n⚠️  [{model_type}] 모델을 찾을 수 없습니다: {model_path}")
                continue
        
        # 테스트 실행
        stats = test_model(model_path, model_type, video_path, 
                          warmup_frames, test_frames, device, show_video=True)
        
        if stats:
            results[model_type] = stats
    
    # 비교 결과 출력
    if len(results) < 2:
        print("\n⚠️  비교할 모델이 충분하지 않습니다 (최소 2개 필요)")
        return
    
    print(f"\n{'='*80}")
    print("📈 최종 비교 결과")
    print(f"{'='*80}")
    
    # 테이블 헤더
    header = f"\n{'항목':<30}"
    for model_type in results.keys():
        header += f"{model_type:<20}"
    print(header)
    print(f"{'-'*80}")
    
    # FPS 비교
    fps_line = f"{'평균 FPS':<30}"
    fps_values = []
    for model_type, stats in results.items():
        fps = stats['avg_fps']
        fps_values.append((model_type, fps))
        fps_line += f"{fps:<20.2f}"
    print(fps_line)
    
    # 추론 시간 비교
    time_line = f"{'평균 추론 시간 (ms)':<30}"
    for model_type, stats in results.items():
        time_ms = stats['avg_inference_time_ms']
        time_line += f"{time_ms:<20.2f}"
    print(time_line)
    
    # 총 처리 시간 비교
    total_line = f"{'총 처리 시간 (초)':<30}"
    for model_type, stats in results.items():
        total_time = stats['total_time']
        total_line += f"{total_time:<20.2f}"
    print(total_line)
    
    # 승자 결정
    print(f"\n{'='*80}")
    best_model = max(fps_values, key=lambda x: x[1])
    print(f"🏆 성능 우승자: {best_model[0]} ({best_model[1]:.2f} FPS)")
    
    # 상대적 성능 비교
    print(f"\n📊 상대적 성능 비교 (기준: {best_model[0]})")
    print(f"{'-'*80}")
    for model_type, fps in fps_values:
        if model_type == best_model[0]:
            print(f"  {model_type:<15}: 100.0% (기준)")
        else:
            relative = (fps / best_model[1]) * 100
            diff = relative - 100
            print(f"  {model_type:<15}: {relative:>5.1f}% ({diff:+.1f}%)")

if __name__ == "__main__":
    compare_models(
        base_path=MODEL_BASE_PATH,
        video_path=VIDEO_PATH,
        device=TEST_DEVICE,
        warmup_frames=WARMUP_FRAMES,
        test_frames=TEST_FRAMES,
        test_models=TEST_MODELS
    )