import cv2 import os import csv import re import numpy as np import pandas as pd import joblib import math import lightgbm import xgboost from util.ears_ai import EarsAI from util.calc_ste_position import CalcStethoscopePosition import torch import torch.nn as nn from torchvision import transforms from PIL import Image import argparse from modules.EARSForDL.model import RegressionResNet # モデル定義をインポート import pickle import matplotlib.pyplot as plt def normalize_quadrilateral_with_point(points, extra_point): all_points = np.vstack([points.reshape(-1, 2), extra_point]) center = np.mean(points.reshape(-1, 2), axis=0) centered_points = all_points - center left_shoulder = centered_points[0] right_shoulder = centered_points[1] shoulder_vector = right_shoulder - left_shoulder angle = np.arctan2(shoulder_vector[1], shoulder_vector[0]) rotation_matrix = np.array([[np.cos(-angle), -np.sin(-angle)], [np.sin(-angle), np.cos(-angle)]]) rotated_points = np.dot(centered_points, rotation_matrix.T) max_edge_length = np.max(np.linalg.norm(np.roll(rotated_points[:4], -1, axis=0) - rotated_points[:4], axis=1)) normalized_points = rotated_points / max_edge_length return normalized_points def normalize_quadrilateral_with_point_average_rotation(points, extra_point): all_points = np.vstack([points.reshape(-1, 2), extra_point]) center = np.mean(points.reshape(-1, 2), axis=0) centered_points = all_points - center left_shoulder, right_shoulder, left_hip, right_hip = centered_points[:4] shoulder_angle = calculate_rotation_angle(left_shoulder, right_shoulder) hip_angle = calculate_rotation_angle(left_hip, right_hip) average_angle = (shoulder_angle + hip_angle) / 2 rotation_matrix = np.array( [[np.cos(-average_angle), -np.sin(-average_angle)], [np.sin(-average_angle), np.cos(-average_angle)]] ) rotated_points = np.dot(centered_points, rotation_matrix.T) max_edge_length = np.max(np.linalg.norm(np.roll(rotated_points[:4], -1, axis=0) - rotated_points[:4], axis=1)) normalized_points = rotated_points / max_edge_length return normalized_points def calculate_rotation_angle(point1, point2): vector = point2 - point1 return np.arctan2(vector[1], vector[0]) def video_to_frames(video_path, output_dir): if not os.path.exists(output_dir): os.makedirs(output_dir) video = cv2.VideoCapture(video_path) if not video.isOpened(): raise IOError(f"動画ファイルを開けませんでした: {video_path}") frame_num = 0 while True: success, frame = video.read() if not success: break frame_num += 1 output_filename = f"{frame_num}-frame.png" cv2.imwrite(os.path.join(output_dir, output_filename), frame) video.release() print(f"全てのフレームを {output_dir} に保存しました。") def lgb_load_model(model_path): with open(model_path, 'rb') as model_file: loaded_model = pickle.load(model_file) return loaded_model def xg_load_model(model_path): with open(model_path, 'rb') as model_file: loaded_model = pickle.load(model_file) return loaded_model def CNN_load_model(model_path, device, resnet_depth=18): model = RegressionResNet(resnet_depth) model.load_state_dict(torch.load(model_path, map_location=device)) model.to(device) model.eval() return model def predict(model, data): return model.predict(data) def calculate_distance(point1, point2): return math.sqrt((point1[0] - point2[0]) ** 2 + (point1[1] - point2[1]) ** 2) def preprocess_image(image_path): transform = transforms.Compose( [ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) image = Image.open(image_path).convert("RGB") return transform(image).unsqueeze(0) def cnn_predict(model, image_tensor, device): with torch.no_grad(): output = model(image_tensor.to(device)) return output.cpu().numpy()[0] def process_images(base_dir, draw_trajectory=True): ears_ai = EarsAI() calc_position = CalcStethoscopePosition() images_dir = base_dir results_dir = os.path.join(os.path.dirname(base_dir), "results") csv_path = os.path.join(results_dir, "results.csv") pose_overlay_dir = os.path.join(results_dir, "pose_overlay_image") stethoscope_overlay_dir = os.path.join(results_dir, "stethoscope_overlay_image") os.makedirs(results_dir, exist_ok=True) os.makedirs(pose_overlay_dir, exist_ok=True) os.makedirs(stethoscope_overlay_dir, exist_ok=True) png_files = [f for f in os.listdir(images_dir) if f.lower().endswith(".png")] png_files.sort(key=lambda x: int(re.search(r"(\d+)", x).group(1))) rows = [] for image_file_name in png_files: image_path = os.path.join(images_dir, image_file_name) print(f"Processing image: {image_path}") frame = cv2.imread(image_path) if frame is None: print(f"Failed to load image: {image_path}") continue pose_overlay_img, left_shoulder, right_shoulder, left_hip, right_hip = ears_ai.pose_detect(frame, None) stethoscope_overlay_img, stethoscope_x, stethoscope_y = ears_ai.ssd_detect(frame, None) pose_overlay_path = os.path.join(pose_overlay_dir, image_file_name) cv2.imwrite(pose_overlay_path, cv2.cvtColor(pose_overlay_img, cv2.COLOR_RGB2BGR)) print(f"Saved pose overlay image: {pose_overlay_path}") stethoscope_overlay_path = os.path.join(stethoscope_overlay_dir, image_file_name) cv2.imwrite(stethoscope_overlay_path, cv2.cvtColor(stethoscope_overlay_img, cv2.COLOR_RGB2BGR)) print(f"Saved stethoscope overlay image: {stethoscope_overlay_path}") # S5とS6の計算 S5 = calculate_distance(right_shoulder, left_hip) S6 = calculate_distance(left_shoulder, right_hip) # S5とS6の比率計算 S5_standard = 215 S6_standard = 204 S5_ratio = S5 / S5_standard S6_ratio = S6 / S6_standard theta_1 = ((-4.5 * 100 * S5_ratio + 440) + (5.0 * 100 * S6_ratio - 500)) / 2 theta_2 = ((53 * S5_ratio - 53) + (57 * S6_ratio - 57)) / 2 x_e = 1.01 * theta_1 + 0.58 y_e = 0.79 * theta_2 - 0.45 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # モデルの読み込み """ cnn_model = CNN_load_model("./models/best_model.pth", device, 18) image_tensor = preprocess_image(image_path) cnn_prediction = cnn_predict(cnn_model, image_tensor, device) cnn_stethoscope_x = int(round(cnn_prediction[0])) cnn_stethoscope_y = int(round(cnn_prediction[1])) """ row = { "image_file_name": image_file_name, "left_shoulder_x": left_shoulder[1], "left_shoulder_y": left_shoulder[0], "right_shoulder_x": right_shoulder[1], "right_shoulder_y": right_shoulder[0], "left_hip_x": left_hip[1], "left_hip_y": left_hip[0], "right_hip_x": right_hip[1], "right_hip_y": right_hip[0], "stethoscope_x": stethoscope_x, "stethoscope_y": stethoscope_y, } """ row = { "image_file_name": image_file_name, "left_shoulder_x": left_shoulder[1], "left_shoulder_y": left_shoulder[0], "right_shoulder_x": right_shoulder[1], "right_shoulder_y": right_shoulder[0], "left_hip_x": left_hip[1], "left_hip_y": left_hip[0], "right_hip_x": right_hip[1], "right_hip_y": right_hip[0], "stethoscope_x": stethoscope_x, "stethoscope_y": stethoscope_y, "cnn_stethoscope_x": cnn_stethoscope_x, "cnn_stethoscope_y": cnn_stethoscope_y, } """ rows.append(row) if rows: """ fieldnames = list(rows[0].keys()) + [ "conv_stethoscope_x", "conv_stethoscope_y", "formula_stethoscope_x", "formula_stethoscope_y", "lightGBM_stethoscope_x", "lightGBM_stethoscope_y", ] """ fieldnames = list(rows[0].keys()) + [ "conv_stethoscope_x", "conv_stethoscope_y", "Xgboost_stethoscope_x", "Xgboost_stethoscope_y", "lightGBM_stethoscope_x", "lightGBM_stethoscope_y", ] # Load models # LightGBMモデルをロード lgb_model_x = lgb_load_model("./models/lgb_stethoscope_calc_x_best_model.pkl") lgb_model_y = lgb_load_model("./models/lgb_stethoscope_calc_y_best_model.pkl") # XGBoostをロード xg_model_x = xg_load_model("./models/xg_stethoscope_calc_x_best_model.pkl") xg_model_y = xg_load_model("./models/xg_stethoscope_calc_y_best_model.pkl") with open(csv_path, "w", newline="") as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writeheader() # 変数の初期化 pre_conv_stethoscope_x = pre_conv_stethoscope_y = 180 # pre_formula_stethoscope_x = pre_formula_stethoscope_y = 180 pre_lightGBM_stethoscope_x = pre_lightGBM_stethoscope_y = 180 pre_xgboost_stethoscope_x = pre_xgboost_stethoscope_y = 180 if stethoscope_x == 0 and stethoscope_y == 0: row["stethoscope_x"] = 320 row["stethoscope_x"] = 240 for row in rows: source_points = np.array( [ [float(row["left_shoulder_x"]), float(row["left_shoulder_y"])], [float(row["right_shoulder_x"]), float(row["right_shoulder_y"])], [float(row["left_hip_x"]), float(row["left_hip_y"])], [float(row["right_hip_x"]), float(row["right_hip_y"])], ], dtype=np.float32, ) stethoscope_x = float(row["stethoscope_x"]) stethoscope_y = float(row["stethoscope_y"]) quadrilateral_points = np.array( [ float(row["left_shoulder_x"]), float(row["left_shoulder_y"]), float(row["right_shoulder_x"]), float(row["right_shoulder_y"]), float(row["left_hip_x"]), float(row["left_hip_y"]), float(row["right_hip_x"]), float(row["right_hip_y"]), ] ) stethoscope_point = np.array([float(row["stethoscope_x"]), float(row["stethoscope_y"])]) normalized_points = normalize_quadrilateral_with_point_average_rotation( quadrilateral_points, stethoscope_point ) if stethoscope_x == 0 and stethoscope_y == 0: row["conv_stethoscope_x"] = pre_conv_stethoscope_x row["conv_stethoscope_y"] = pre_conv_stethoscope_y # row["formula_stethoscope_x"] = pre_formula_stethoscope_x # row["formula_stethoscope_y"] = pre_formula_stethoscope_y row["lightGBM_stethoscope_x"] = pre_lightGBM_stethoscope_x row["lightGBM_stethoscope_y"] = pre_lightGBM_stethoscope_y row["Xgboost_stethoscope_x"] = pre_xgboost_stethoscope_x row["Xgboost_stethoscope_y"] = pre_xgboost_stethoscope_y else: conv_stethoscope = calc_position.calc_affine(source_points, stethoscope_x, stethoscope_y) row["conv_stethoscope_x"], row["conv_stethoscope_y"] = conv_stethoscope """ row["formula_stethoscope_x"], row["formula_stethoscope_y"] = int(conv_stethoscope[0] - x_e), int( conv_stethoscope[1] - y_e ) """ row_convert = { "left_shoulder_x": normalized_points[0, 0], "left_shoulder_y": normalized_points[0, 1], "right_shoulder_x": normalized_points[1, 0], "right_shoulder_y": normalized_points[1, 1], "left_hip_x": normalized_points[2, 0], "left_hip_y": normalized_points[2, 1], "right_hip_x": normalized_points[3, 0], "right_hip_y": normalized_points[3, 1], "stethoscope_x": normalized_points[4, 0], "stethoscope_y": normalized_points[4, 1], } # 各点をプロット plt.scatter(row_convert['left_shoulder_x'], row_convert['left_shoulder_y'], color='blue', s=100, label='Left Shoulder') plt.scatter(row_convert['right_shoulder_x'], row_convert['right_shoulder_y'], color='blue', s=100, label='Right Shoulder') plt.scatter(row_convert['left_hip_x'], row_convert['left_hip_y'], color='green', s=100, label='Left Hip') plt.scatter(row_convert['right_hip_x'], row_convert['right_hip_y'], color='green', s=100, label='Right Hip') plt.scatter(row_convert['stethoscope_x'], row_convert['stethoscope_y'], color='red', s=100, label='Stethoscope') # Machine Learning prediction input_columns = [ "left_shoulder_x", "left_shoulder_y", "right_shoulder_x", "right_shoulder_y", "left_hip_x", "left_hip_y", "right_hip_x", "right_hip_y", "stethoscope_x", "stethoscope_y", ] input_data = pd.DataFrame([row_convert]) row["lightGBM_stethoscope_x"] = int(predict(lgb_model_x, input_data[input_columns])[0]) row["lightGBM_stethoscope_y"] = int(predict(lgb_model_y, input_data[input_columns])[0]) row["Xgboost_stethoscope_x"] = int(predict(xg_model_x, input_data[input_columns])[0]) row["Xgboost_stethoscope_y"] = int(predict(xg_model_y, input_data[input_columns])[0]) pre_conv_stethoscope_x = row["conv_stethoscope_x"] pre_conv_stethoscope_y = row["conv_stethoscope_y"] # pre_formula_stethoscope_x = row["formula_stethoscope_x"] # pre_formula_stethoscope_y = row["formula_stethoscope_y"] pre_lightGBM_stethoscope_x = row["lightGBM_stethoscope_x"] pre_lightGBM_stethoscope_y = row["lightGBM_stethoscope_y"] pre_xgboost_stethoscope_x = row["Xgboost_stethoscope_x"] pre_xgboost_stethoscope_y = row["Xgboost_stethoscope_y"] writer.writerow(row) print(f"Processed and saved results to: {csv_path}") # Add these new variables at the beginning of your script conv_points = [] cnn_points = [] lgbm_points = [] # 新しい処理を追加 df = pd.read_csv(csv_path) original_images_dir = base_dir body_image = cv2.imread("images/body/BodyF.png") results_dir = "images/body/results" os.makedirs(results_dir, exist_ok=True) marked_images_dir = os.path.join(results_dir, "marked_images") os.makedirs(marked_images_dir, exist_ok=True) # 新しいディレクトリを作成 conv_dir = os.path.join(results_dir, "conv") cnn_dir = os.path.join(results_dir, "cnn") lgbm_dir = os.path.join(results_dir, "lgbm") os.makedirs(conv_dir, exist_ok=True) os.makedirs(cnn_dir, exist_ok=True) os.makedirs(lgbm_dir, exist_ok=True) for _, row in df.iterrows(): image_file_name = row["image_file_name"] conv_x, conv_y = int(row["conv_stethoscope_x"]), int(row["conv_stethoscope_y"]) cnn_x, cnn_y = int(row["Xgboost_stethoscope_x"]), int(row["Xgboost_stethoscope_x"]) lgbm_x, lgbm_y = int(row["lightGBM_stethoscope_x"]), int(row["lightGBM_stethoscope_y"]) original_image_path = os.path.join(original_images_dir, image_file_name) original_image = cv2.imread(original_image_path) if original_image is None: print(f"Failed to load image: {original_image_path}") continue # Draw markers in cyan color (BGR: 255, 255, 0) cv2.circle( original_image, (int(row["left_shoulder_x"]), int(row["left_shoulder_y"])), 10, (255, 255, 0), -1 ) cv2.circle( original_image, (int(row["right_shoulder_x"]), int(row["right_shoulder_y"])), 10, (255, 255, 0), -1 ) cv2.circle(original_image, (int(row["left_hip_x"]), int(row["left_hip_y"])), 10, (255, 255, 0), -1) cv2.circle(original_image, (int(row["right_hip_x"]), int(row["right_hip_y"])), 10, (255, 255, 0), -1) cv2.circle(original_image, (int(row["stethoscope_x"]), int(row["stethoscope_y"])), 10, (255, 255, 0), -1) # Save marked image marked_image_path = os.path.join(marked_images_dir, image_file_name) cv2.imwrite(marked_image_path, original_image) conv_points.append((conv_x, conv_y)) cnn_points.append((cnn_x, cnn_y)) lgbm_points.append((lgbm_x, lgbm_y)) # Conv画像を生成 conv_image = body_image.copy() if draw_trajectory and len(conv_points) > 1: cv2.polylines(conv_image, [np.array(conv_points)], False, (0, 255, 0), 2) cv2.circle(conv_image, (conv_x, conv_y), 10, (0, 255, 0), -1) cv2.imwrite(os.path.join(conv_dir, image_file_name), conv_image) # CNN画像を生成 cnn_image = body_image.copy() if draw_trajectory and len(cnn_points) > 1: cv2.polylines(cnn_image, [np.array(cnn_points)], False, (255, 0, 0), 2) cv2.circle(cnn_image, (cnn_x, cnn_y), 10, (255, 0, 0), -1) cv2.imwrite(os.path.join(cnn_dir, image_file_name), cnn_image) # LightGBM画像を生成 lgbm_image = body_image.copy() if draw_trajectory and len(lgbm_points) > 1: cv2.polylines(lgbm_image, [np.array(lgbm_points)], False, (0, 0, 255), 2) cv2.circle(lgbm_image, (lgbm_x, lgbm_y), 10, (0, 0, 255), -1) cv2.imwrite(os.path.join(lgbm_dir, image_file_name), lgbm_image) # 動画を生成 create_video_from_images(conv_dir, os.path.join(results_dir, "conv_video_with_trajectory.mp4")) create_video_from_images(cnn_dir, os.path.join(results_dir, "cnn_video_with_trajectory.mp4")) create_video_from_images(lgbm_dir, os.path.join(results_dir, "lgbm_video_with_trajectory.mp4")) # 軌跡なしの動画を生成 create_video_from_images(conv_dir, os.path.join(results_dir, "conv_video_without_trajectory.mp4"), False) create_video_from_images(cnn_dir, os.path.join(results_dir, "cnn_video_without_trajectory.mp4"), False) create_video_from_images(lgbm_dir, os.path.join(results_dir, "lgbm_video_without_trajectory.mp4"), False) # Create video from marked images create_video_from_images(marked_images_dir, os.path.join(results_dir, "marked_video.mp4")) else: print("No data to write to CSV.") def create_video_from_images(image_dir, output_path, with_trajectory=True): images = [img for img in os.listdir(image_dir) if img.endswith(".png")] images.sort(key=lambda x: int(re.search(r"(\d+)", x).group())) if images: frame = cv2.imread(os.path.join(image_dir, images[0])) height, width, layers = frame.shape video = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), 30, (width, height)) for image in images: img = cv2.imread(os.path.join(image_dir, image)) if not with_trajectory: # 軌跡を消去（背景画像で上書き） background = cv2.imread("images/body/BodyF.png") mask = cv2.threshold(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 10, 255, cv2.THRESH_BINARY)[1] img = cv2.bitwise_and(img, img, mask=mask) background = cv2.bitwise_and(background, background, mask=cv2.bitwise_not(mask)) img = cv2.add(img, background) video.write(img) cv2.destroyAllWindows() video.release() print(f"Created video: {output_path}") else: print(f"No images found in {image_dir}") if __name__ == "__main__": parser = argparse.ArgumentParser(description="Process video and generate results.") parser.add_argument("--video_path", default="./video/Test3-1.mp4", help="Path to the input video file") parser.add_argument("--output_dir", default="images", help="Directory to save output images and results") parser.add_argument("--draw_trajectory", action="store_true", help="Draw trajectory in the output video") args = parser.parse_args() # Step 1: Convert video to frames video_to_frames(args.video_path, args.output_dir) # Step 2: Process the generated images process_images(args.output_dir, args.draw_trajectory)