1+ """minifig向き検出クラス.
2+
3+ @file minifig_detector.py
4+ @author Hara1274
5+ """
6+ import cv2
7+ import numpy as np
8+ import onnxruntime as ort
9+ import os
10+ import json
11+
12+
13+
14+
15+ class MinifigDetector :
16+ """minifigの向きを検出するクラス."""
17+
18+ def __init__ (self ):
19+ """検出器を初期化."""
20+ self .session = None
21+ self .input_name = None
22+ self .labels = ["front" , "back" , "right" , "left" ]
23+ self .conf_threshold = 0.25
24+ self .nms_threshold = 0.45
25+ self .input_size = 640
26+
27+ # YOLO11 ONNXモデルのパス
28+ project_root = os .path .dirname (os .path .dirname (__file__ ))
29+ model_path = os .path .join (project_root , "models" , "11s_100epoch_&_650imgsz_fig.onnx" )
30+
31+ # モデルファイルが存在する場合のみonnxモデルの設定
32+ if os .path .exists (model_path ):
33+ # 推論セッションを作成
34+ self .session = ort .InferenceSession (model_path )
35+ # ONNXモデルの入力層の名前を取得
36+ self .input_name = self .session .get_inputs ()[0 ].name
37+
38+ def preprocess_image (self , img ):
39+ """推論用に画像を前処理.
40+
41+ Args:
42+ img: 入力画像
43+
44+ Returns:
45+ tuple: (処理後画像, スケール比, パディング情報)
46+ """
47+ # アスペクト比を保持するスケール計算
48+ shape = img .shape [:2 ] # 元画像サイズ (H, W)
49+ r = min (self .input_size / shape [0 ], self .input_size / shape [1 ]) # アスペクト比維持のスケール
50+ new_unpad = (int (round (shape [1 ] * r )), int (round (shape [0 ] * r ))) # スケール後サイズ (W, H)
51+ img_resized = cv2 .resize (img , new_unpad , interpolation = cv2 .INTER_LINEAR )
52+
53+ # 640x640にするための余白計算
54+ dw = self .input_size - new_unpad [0 ] # 水平余白
55+ dh = self .input_size - new_unpad [1 ] # 垂直余白
56+ top , bottom = dh // 2 , dh - dh // 2 # 上下分割
57+ left , right = dw // 2 , dw - dw // 2 # 左右分割
58+
59+ # 灰色（114）でパディングして640x640に調整
60+ img_padded = cv2 .copyMakeBorder (img_resized , top , bottom , left , right , cv2 .BORDER_CONSTANT , value = (114 , 114 , 114 ))
61+
62+ # YOLO入力形式に変換: HWC→CHW, [0,1]正規化, バッチ次元追加
63+ img_input = img_padded .transpose (2 , 0 , 1 ).astype (np .float32 ) / 255.0 # (3,640,640)
64+ img_input = np .expand_dims (img_input , axis = 0 ) # (1,3,640,640)
65+
66+ return img_input , r , (left , top )
67+
68+ def postprocess (self , pred , scale , pad ):
69+ """推論結果を後処理.
70+
71+ Args:
72+ pred: 推論出力 [1, 8, N] (YOLO11形式)
73+ scale: レターボックスのスケール比
74+ pad: パディング情報 (left, top)
75+
76+ Returns:
77+ str: 検出された向き
78+ """
79+ if pred .size == 0 :
80+ return "front"
81+
82+ # YOLO11出力解析: [1,8,N] -> [8,N]（バッチ次元除去）
83+ data = pred [0 ]
84+
85+ attributes = 8 # [cx,cy,w,h,class0,class1,class2,class3]
86+ num_classes = attributes - 4 # 座標4つ以外はクラススコア
87+ num_boxes = data .shape [1 ] # 検出候補数を取得
88+
89+ boxes = []
90+ confidences = []
91+ class_ids = []
92+
93+ # 各検出候補を処理
94+ for i in range (num_boxes ):
95+ max_score = - 1.0
96+ best_class = - 1
97+
98+ # クラススコアの最大値とクラスIDを取得
99+ for j in range (num_classes ):
100+ score = data [4 + j , i ]
101+ if score > max_score :
102+ max_score = score
103+ best_class = j
104+
105+ # 信頼度閾値による候補フィルタリング
106+ if max_score < self .conf_threshold :
107+ continue
108+
109+ # バウンディングボックス座標抽出（YOLO形式：中心座標＋幅高さ）
110+ cx = data [0 , i ] # 中心X座標（640x640空間）
111+ cy = data [1 , i ] # 中心Y座標（640x640空間）
112+ w = data [2 , i ] # 幅（640x640空間）
113+ h = data [3 , i ] # 高さ（640x640空間）
114+
115+ # レターボックス座標系から元画像座標系に逆変換
116+ center_x = int ((cx - pad [0 ]) / scale ) # パディング補正＋スケール逆変換
117+ center_y = int ((cy - pad [1 ]) / scale )
118+ width = int (w / scale )
119+ height = int (h / scale )
120+
121+ # 中心座標形式から左上座標形式（OpenCV形式）に変換
122+ left = center_x - width // 2
123+ top = center_y - height // 2
124+
125+ # NMS用データに追加
126+ boxes .append ([left , top , width , height ]) # [x,y,w,h]形式
127+ confidences .append (max_score ) # 最高クラススコア
128+ class_ids .append (best_class ) # 最高スコアのクラスID
129+
130+ # 信頼度閾値を満たす検出がない場合
131+ if not boxes :
132+ return {"wasDetected" : False , "direction" : "NONE" , "confidence" : 0.0 }
133+
134+ # Non-Maximum Suppression で重複検出を除去
135+ indices = cv2 .dnn .NMSBoxes (boxes , confidences , self .conf_threshold , self .nms_threshold )
136+
137+ # NMS後に有効な検出が残っている場合
138+ if len (indices ) > 0 :
139+ best_idx = indices .flatten ()[0 ] # 最も信頼度の高い検出を選択
140+ best_class_id = class_ids [best_idx ] # 対応するクラスID
141+ best_confidence = confidences [best_idx ] # 対応する信頼度
142+ direction = self .labels [best_class_id ] #ミニフィグの向き
143+
144+ return {
145+ "wasDetected" : True ,
146+ "direction" : direction , # "front", "back", "right", "left"
147+ "confidence" : float (best_confidence ) # 0.0-1.0の信頼度
148+ }
149+
150+ return {"wasDetected" : False , "direction" : "NONE" , "confidence" : 0.0 }
151+
152+ def detect (self , image_path : str ) -> dict :
153+ """minifigの向きを判定.
154+
155+ Args:
156+ image_path (str): 画像ファイルのパス
157+
158+ Returns:
159+ dict: 検出結果 {"wasDetected": bool, "direction": str, "confidence": float}
160+ """
161+ # 共通エラーレスポンス
162+ error_result = {"wasDetected" : False , "direction" : "NONE" , "confidence" : 0.0 }
163+
164+ # 入力ファイル存在チェック
165+ if not os .path .exists (image_path ):
166+ return error_result
167+
168+ # モデル読み込み状態チェック
169+ if self .session is None :
170+ return error_result
171+
172+ # 画像読み込み（BGR形式）
173+ img = cv2 .imread (image_path )
174+ if img is None : # 読み込み失敗（不正ファイル等）
175+ return error_result
176+
177+ # 前処理：レターボックス＋正規化
178+ img_input , scale , pad = self .preprocess_image (img )
179+
180+ # YOLO11推論実行
181+ outputs = self .session .run (None , {self .input_name : img_input })
182+ pred = outputs [0 ] # メイン出力 [1,8,N]
183+
184+ # 後処理：NMS＋結果構築
185+ result = self .postprocess (pred , scale , pad )
186+
187+ return result
0 commit comments