fix: 헤드 ONNX 출력 순서 수정 (v0.2.1) [AI]

convert_to_onnx.py의 헤드 변환 시 출력 순서를 정렬하여
원본 PyTorch 모델과 순서가 불일치하던 문제를 수정했습니다.

변경 사항:
- MultiHeadWrapper.forward(): 정렬 제거, 원본 ModuleDict 순서 유지
- output_names: sorted() 제거, list(model.heads.keys()) 사용
- 헤드 순서: ['scene', 'concept', 'object']로 통일

이제 분류 결과가 올바른 헤드에 매핑됩니다.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

head.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:acf595e6f5208e5ebd833868b62a8eb513d56a681b160231f2f10bf25d4b48d1
+oid sha256:ead5678e0ec6c628cbcdd946258d9286ddb53d8969f9a2e3e39c1d5dc718c2ec
 size 90957

inference_example.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 """
 ONNX 모델을 사용한 멀티헤드 이미지 분류 추론 예제
+전체 모델(model.onnx) 또는 분리 모델(encoder.onnx + head.onnx) 사용 가능
 """
 
 import onnxruntime as ort
@@ -8,6 +9,7 @@ import numpy as np
 from PIL import Image
 import torchvision.transforms as transforms
 import json
+from pathlib import Path
 
 # 전처리 파이프라인
 transform = transforms.Compose([
@@ -34,14 +36,14 @@ def softmax(x):
     exp_x = np.exp(x - np.max(x, axis=1, keepdims=True))
     return exp_x / np.sum(exp_x, axis=1, keepdims=True)
 
-def predict_image(onnx_model_path, model_info_path, image_path):
-    """이미지 분류 예측"""
+def predict_image_full_model(model_path, model_info_path, image_path):
+    """전체 모델을 사용한 이미지 분류 예측"""
 
     # 모델 정보 로드
     model_info = load_model_info(model_info_path)
 
     # ONNX 세션 생성
-    session = ort.InferenceSession(onnx_model_path)
+    session = ort.InferenceSession(model_path)
 
     # 이미지 전처리
     image_array = preprocess_image(image_path)
@@ -53,71 +55,111 @@ def predict_image(onnx_model_path, model_info_path, image_path):
     # 결과 해석
     results = {}
     head_names = list(model_info['output_specification']['heads'].keys())
-    output_names = head_names + ['features']  # features 추가
-
-    for i, output_name in enumerate(output_names):
-        if output_name == 'features':
-            # 특징 벡터 처리
-            features = outputs[i][0]  # 첫 번째 배치
-            results[output_name] = {
-                'embedding': features.tolist(),
-                'dimension': len(features),
-                'description': 'DINOv2 backbone features'
-            }
-        else:
-            # 분류 헤드 처리
-            logits = outputs[i]
-            probabilities = softmax(logits)[0]  # 첫 번째 배치
-
-            # 클래스 이름 매핑
-            class_names = model_info['class_mappings'].get(output_name, {})
-
-            # 최고 확률 클래스
-            pred_idx = np.argmax(probabilities)
-            pred_class = class_names.get(str(pred_idx), f"Class_{pred_idx}")
-            pred_prob = probabilities[pred_idx]
-
-            # 상위 3개 클래스
-            top3_indices = np.argsort(probabilities)[-3:][::-1]
-            top3_results = []
-            for idx in top3_indices:
-                class_name = class_names.get(str(idx), f"Class_{idx}")
-                prob = probabilities[idx]
-                top3_results.append({'class': class_name, 'probability': float(prob)})
-
-            results[output_name] = {
-                'predicted_class': pred_class,
-                'confidence': float(pred_prob),
-                'top3': top3_results
-            }
+
+    for i, output_name in enumerate(head_names):
+        logits = outputs[i]
+        probabilities = softmax(logits)[0]
+
+        # 클래스 이름 매핑
+        class_names = model_info['class_mappings'].get(output_name, {})
+
+        # 최고 확률 클래스
+        pred_idx = np.argmax(probabilities)
+        pred_class = class_names.get(str(pred_idx), f"Class_{pred_idx}")
+        pred_prob = probabilities[pred_idx]
+
+        # 상위 3개 클래스
+        top3_indices = np.argsort(probabilities)[-3:][::-1]
+        top3_results = []
+        for idx in top3_indices:
+            class_name = class_names.get(str(idx), f"Class_{idx}")
+            prob = probabilities[idx]
+            top3_results.append({'class': class_name, 'probability': float(prob)})
+
+        results[output_name] = {
+            'predicted_class': pred_class,
+            'confidence': float(pred_prob),
+            'top3': top3_results
+        }
+
+    return results
+
+def predict_image_split_model(encoder_path, head_path, model_info_path, image_path):
+    """분리 모델을 사용한 이미지 분류 예측"""
+
+    # 모델 정보 로드
+    model_info = load_model_info(model_info_path)
+
+    # ONNX 세션 생성
+    encoder_session = ort.InferenceSession(encoder_path)
+    head_session = ort.InferenceSession(head_path)
+
+    # 이미지 전처리
+    image_array = preprocess_image(image_path)
+
+    # 인코더로 특징 벡터 추출
+    encoder_inputs = {'image': image_array}
+    features = encoder_session.run(None, encoder_inputs)[0]
+
+    # 헤드로 분류
+    head_inputs = {'features': features}
+    outputs = head_session.run(None, head_inputs)
+
+    # 결과 해석
+    results = {}
+    head_names = list(model_info['output_specification']['heads'].keys())
+
+    for i, output_name in enumerate(head_names):
+        logits = outputs[i]
+        probabilities = softmax(logits)[0]
+
+        # 클래스 이름 매핑
+        class_names = model_info['class_mappings'].get(output_name, {})
+
+        # 최고 확률 클래스
+        pred_idx = np.argmax(probabilities)
+        pred_class = class_names.get(str(pred_idx), f"Class_{pred_idx}")
+        pred_prob = probabilities[pred_idx]
+
+        # 상위 3개 클래스
+        top3_indices = np.argsort(probabilities)[-3:][::-1]
+        top3_results = []
+        for idx in top3_indices:
+            class_name = class_names.get(str(idx), f"Class_{idx}")
+            prob = probabilities[idx]
+            top3_results.append({'class': class_name, 'probability': float(prob)})
+
+        results[output_name] = {
+            'predicted_class': pred_class,
+            'confidence': float(pred_prob),
+            'top3': top3_results
+        }
 
     return results
 
 # 사용 예시
 if __name__ == "__main__":
-    onnx_path = "image_classifier.onnx"
     model_info_path = "model_info.json"
     image_path = "test_image.jpg"
 
-    try:
-        results = predict_image(onnx_path, model_info_path, image_path)
-
-        print(f"이미지 분류 결과: {image_path}")
-        print("=" * 50)
-
-        for output_name, result in results.items():
-            if output_name == 'features':
-                print(f"\n{output_name.upper()}:")
-                print(f"  차원: {result['dimension']}")
-                print(f"  설명: {result['description']}")
-                print(f"  특징 벡터 (처음 10개): {result['embedding'][:10]}")
-            else:
-                print(f"\n{output_name.upper()}:")
-                print(f"  예측 클래스: {result['predicted_class']}")
-                print(f"  신뢰도: {result['confidence']:.4f}")
-                print(f"  Top 3:")
-                for i, top_result in enumerate(result['top3'], 1):
-                    print(f"    {i}. {top_result['class']}: {top_result['probability']:.4f}")
-
-    except Exception as e:
-        print(f"추론 실패: {e}")
+    # 분리 모델이 있는지 확인
+    if Path("encoder.onnx").exists() and Path("head.onnx").exists():
+        print("분리 모델 사용")
+        results = predict_image_split_model("encoder.onnx", "head.onnx", model_info_path, image_path)
+    elif Path("model.onnx").exists():
+        print("전체 모델 사용")
+        results = predict_image_full_model("model.onnx", model_info_path, image_path)
+    else:
+        print("ONNX 모델을 찾을 수 없습니다.")
+        exit(1)
+
+    print(f"\n이미지 분류 결과: {image_path}")
+    print("=" * 50)
+
+    for output_name, result in results.items():
+        print(f"\n{output_name.upper()}:")
+        print(f"  예측 클래스: {result['predicted_class']}")
+        print(f"  신뢰도: {result['confidence']:.4f}")
+        print(f"  Top 3:")
+        for i, top_result in enumerate(result['top3'], 1):
+            print(f"    {i}. {top_result['class']}: {top_result['probability']:.4f}")