diff --git a/main.py b/main.py index c7993bf..d492f40 100644 --- a/main.py +++ b/main.py @@ -60,28 +60,95 @@ return None -def yolox_detector_inference(frame, yolox_inferencer, score_thr=0.3): - # yolox_inferencerがNoneの場合、デフォルト値を返す - if yolox_inferencer is None: - return frame, 0, 0 +def is_point_in_polygon(point, polygon): + """ + Ray-Castingアルゴリズムを使用して、点が多角形の内部にあるかどうかを判定する - # 以下は既存の処理 + Args: + point: テストする点 [x, y] + polygon: 多角形の頂点リスト [[x1, y1], [x2, y2], ...] + + Returns: + bool: 点が多角形の内部にある場合はTrue + """ + x, y = point + n = len(polygon) + inside = False + + j = n - 1 + for i in range(n): + if ((polygon[i][1] > y) != (polygon[j][1] > y)) and ( + x + < (polygon[j][0] - polygon[i][0]) + * (y - polygon[i][1]) + / (polygon[j][1] - polygon[i][1]) + + polygon[i][0] + ): + inside = not inside + j = i + + return inside + + +def yolox_detector_inference(frame, yolox_inferencer, landmarks=None, score_thr=0.3): + """ + YOLOXを使用して聴診器を検出し、ランドマークで形成されるポリゴン内の座標を返す + + Args: + frame: 入力画像 + yolox_inferencer: YOLOXの推論モデル + landmarks: [[left_shoulder_x, left_shoulder_y], + [right_shoulder_x, right_shoulder_y], + [left_hip_x, left_hip_y], + [right_hip_x, right_hip_y]] + score_thr: 検出スコアの閾値 + + Returns: + tuple: (オーバーレイ画像, 検出されたx座標, 検出されたy座標) + """ + # BGRからRGBに変換 frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # 検出を実行 result = yolox_inferencer(inputs=frame_rgb, return_vis=True) + + # 検出結果を取得 predictions = result["predictions"][0] stethoscope_x = None stethoscope_y = None max_score = -1 + # ランドマークが与えられている場合、ポリゴンを形成 + if landmarks is not None: + polygon = landmarks + else: + # ランドマークがない場合は画像全体を対象とする + h, w = frame.shape[:2] + polygon = [[0, 0], [w, 0], [w, h], [0, h]] + for i, (label, score) in enumerate( zip(predictions["labels"], predictions["scores"]) ): - if score >= score_thr and label == 0 and score > max_score: + if score >= score_thr and label == 0: # label 0 は聴診器を示す bbox = predictions["bboxes"][i] - stethoscope_x = (bbox[0] + bbox[2]) / 2 - stethoscope_y = (bbox[1] + bbox[3]) / 2 - max_score = score + # バウンディングボックスの中心点を計算 + center_x = (bbox[0] + bbox[2]) / 2 + center_y = (bbox[1] + bbox[3]) / 2 + + # 中心点がポリゴン内にあるかチェック + if is_point_in_polygon([center_x, center_y], polygon): + if score > max_score: + stethoscope_x = center_x + stethoscope_y = center_y + max_score = score + + # 検出結果がない場合は0,0を返す + if stethoscope_x is None or stethoscope_y is None: + stethoscope_x = 0 + stethoscope_y = 0 + + # 可視化結果の画像を取得してBGRに変換 stethoscope_overlay_img = result["visualization"][0] if ( len(stethoscope_overlay_img.shape) == 3 @@ -91,11 +158,6 @@ stethoscope_overlay_img, cv2.COLOR_RGB2BGR ) - # 検出結果がない場合は0,0を返す - if stethoscope_x is None or stethoscope_y is None: - stethoscope_x = 0 - stethoscope_y = 0 - return stethoscope_overlay_img, stethoscope_x, stethoscope_y @@ -266,8 +328,14 @@ yolox_inferencer = init_yolox() if YOLOX_ENABLED: + landmarks = [ + [left_shoulder[0], left_shoulder[1]], + [right_shoulder[0], right_shoulder[1]], + [left_hip[0], left_hip[1]], + [right_hip[0], right_hip[1]], + ] stethoscope_overlay_img, stethoscope_x, stethoscope_y = ( - yolox_detector_inference(frame, yolox_inferencer) + yolox_detector_inference(frame, yolox_inferencer, landmarks=landmarks) ) cv2.imwrite(