diff --git a/config.py b/config.py index f531564..c6ad6fe 100644 --- a/config.py +++ b/config.py @@ -3,11 +3,11 @@ YOLOX_CHECKPOINT_FILE = "models/yolox/yolox_x_White_only.pth" DEVICE = "cuda:0" -# Colors for different models (R,G,B format) -CONV_COLOR = (0, 255, 0) # Green -XGBOOST_COLOR = (255, 0, 0) # Red -LIGHTGBM_COLOR = (0, 0, 255) # Blue -EARSNET_COLOR = (0, 165, 255) # Orange +# Colors for different models (B,G,R format) +CONV_COLOR = (0, 255, 0) +XGBOOST_COLOR = (98, 72, 16) +LIGHTGBM_COLOR = (98, 72, 16) +EARSNET_COLOR = (18, 53, 13) # Model execution settings CONV_ENABLED = True diff --git a/main.py b/main.py index 3a7a025..62173aa 100644 --- a/main.py +++ b/main.py @@ -331,6 +331,7 @@ # process_images関数内のRTMPOSE_ENABLEDの部分を修正 if RTMPOSE_ENABLED: + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) det_result = inference_detector(detector, frame) pred_instance = det_result.pred_instances.cpu().numpy() bboxes = np.concatenate( @@ -386,14 +387,25 @@ ) ) - cv2.imwrite( - os.path.join(pose_overlay_dir, image_file_name), - cv2.cvtColor(pose_overlay_img, cv2.COLOR_RGB2BGR), - ) - cv2.imwrite( - os.path.join(stethoscope_overlay_dir, image_file_name), - cv2.cvtColor(stethoscope_overlay_img, cv2.COLOR_RGB2BGR), - ) + if RTMPOSE_ENABLED and YOLOX_ENABLED: + # RTMPoseとYOLOXの場合は既にBGRなので変換不要 + cv2.imwrite( + os.path.join(pose_overlay_dir, image_file_name), pose_overlay_img + ) + cv2.imwrite( + os.path.join(stethoscope_overlay_dir, image_file_name), + stethoscope_overlay_img, + ) + else: + # その他の場合は従来通り色変換を行う + cv2.imwrite( + os.path.join(pose_overlay_dir, image_file_name), + cv2.cvtColor(pose_overlay_img, cv2.COLOR_RGB2BGR), + ) + cv2.imwrite( + os.path.join(stethoscope_overlay_dir, image_file_name), + cv2.cvtColor(stethoscope_overlay_img, cv2.COLOR_RGB2BGR), + ) if POSENET_ENABLED: row = {