diff --git a/main.py b/main.py index b030988..d5bf503 100644 --- a/main.py +++ b/main.py @@ -28,6 +28,9 @@ XGBOOST_ENABLED = os.getenv("XGBOOST_ENABLED", "True").lower() == "true" LIGHTGBM_ENABLED = os.getenv("LIGHTGBM_ENABLED", "True").lower() == "true" +# Get normalization setting +NORMALIZE_ENABLED = os.getenv("NORMALIZE_ENABLED", "False").lower() == "true" + def load_model(model_path, model_type="lgb"): with open(model_path, "rb") as model_file: @@ -104,6 +107,7 @@ ) rows = [] + normalized_rows = [] for image_file_name in png_files: image_path = os.path.join(base_dir, image_file_name) frame = cv2.imread(image_path) @@ -135,6 +139,32 @@ } rows.append(row) + if NORMALIZE_ENABLED: + source_points = np.array( + [ + [float(row[f"{pos}_x"]), float(row[f"{pos}_y"])] + for pos in ["left_shoulder", "right_shoulder", "left_hip", "right_hip"] + ], + dtype=np.float32, + ) + stethoscope_point = np.array([float(row["stethoscope_x"]), float(row["stethoscope_y"])]) + normalized_points = normalize_quadrilateral_with_point(source_points.flatten(), stethoscope_point) + + normalized_row = { + "image_file_name": image_file_name, + "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], + } + normalized_rows.append(normalized_row) + if rows: fieldnames = list(rows[0].keys()) if CONV_ENABLED: @@ -163,7 +193,7 @@ if XGBOOST_ENABLED: prev_values["Xgboost"] = (180, 180) - for row in rows: + for i, row in enumerate(rows): source_points = np.array( [ [float(row[f"{pos}_x"]), float(row[f"{pos}_y"])] @@ -204,9 +234,28 @@ for key in prev_values: prev_values[key] = (row[f"{key}_stethoscope_x"], row[f"{key}_stethoscope_y"]) + # 正規化されたデータにも予測結果を追加 + if NORMALIZE_ENABLED: + for key in prev_values: + normalized_rows[i][f"{key}_stethoscope_x"] = row[f"{key}_stethoscope_x"] + normalized_rows[i][f"{key}_stethoscope_y"] = row[f"{key}_stethoscope_y"] + writer.writerow(row) print(f"Processed and saved results to: {csv_path}") + + # Write normalized results-convert.csv if enabled + if NORMALIZE_ENABLED and normalized_rows: + normalized_csv_path = os.path.join(results_dir, "results-convert.csv") + normalized_fieldnames = list(normalized_rows[0].keys()) + + with open(normalized_csv_path, "w", newline="") as csvfile: + writer = csv.DictWriter(csvfile, fieldnames=normalized_fieldnames) + writer.writeheader() + writer.writerows(normalized_rows) + + print(f"Processed and saved normalized results to: {normalized_csv_path}") + generate_visualizations(csv_path, base_dir, output_dir) else: print("No data to write to CSV.") @@ -214,7 +263,7 @@ def generate_visualizations(csv_path, original_images_dir, output_dir): df = pd.read_csv(csv_path) - body_image = cv2.imread("./images/body/BodyF.png") # オリジナルのパスを使用 + body_image = cv2.imread("./images/body/BodyF.png") results_dir = os.path.join(output_dir, "results") os.makedirs(results_dir, exist_ok=True) @@ -320,4 +369,3 @@ frames_dir = os.path.join(args.output_dir, "frames") video_to_frames(args.video_path, frames_dir) process_images(frames_dir, args.output_dir) - generate_visualizations(os.path.join(args.output_dir, "results", "results.csv"), frames_dir, args.output_dir)