diff --git a/main.py b/main.py index d5bf503..25240f6 100644 --- a/main.py +++ b/main.py @@ -94,6 +94,7 @@ calc_position = CalcStethoscopePosition() results_dir = os.path.join(output_dir, "results") csv_path = os.path.join(results_dir, "results.csv") + normalized_csv_path = os.path.join(results_dir, "results-convert.csv") pose_overlay_dir = os.path.join(results_dir, "pose_overlay_image") stethoscope_overlay_dir = os.path.join(results_dir, "stethoscope_overlay_image") @@ -139,31 +140,30 @@ } 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) + 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) + 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()) @@ -181,9 +181,11 @@ xg_model_x = load_model("./models/xg_stethoscope_calc_x_best_model-Fold4.pkl") xg_model_y = load_model("./models/xg_stethoscope_calc_y_best_model-Fold4.pkl") - with open(csv_path, "w", newline="") as csvfile: + with open(csv_path, "w", newline="") as csvfile, open(normalized_csv_path, "w", newline="") as norm_csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + norm_writer = csv.DictWriter(norm_csvfile, fieldnames=fieldnames) writer.writeheader() + norm_writer.writeheader() prev_values = {} if CONV_ENABLED: @@ -193,7 +195,7 @@ if XGBOOST_ENABLED: prev_values["Xgboost"] = (180, 180) - for i, row in enumerate(rows): + for i, (row, norm_row) in enumerate(zip(rows, normalized_rows)): source_points = np.array( [ [float(row[f"{pos}_x"]), float(row[f"{pos}_y"])] @@ -206,55 +208,44 @@ if stethoscope_point[0] == 0 and stethoscope_point[1] == 0: for key in prev_values: row[f"{key}_stethoscope_x"], row[f"{key}_stethoscope_y"] = prev_values[key] + norm_row[f"{key}_stethoscope_x"], norm_row[f"{key}_stethoscope_y"] = prev_values[key] else: if CONV_ENABLED: conv_stethoscope = calc_position.calc_affine(source_points, *stethoscope_point) row["conv_stethoscope_x"], row["conv_stethoscope_y"] = conv_stethoscope + norm_row["conv_stethoscope_x"], norm_row["conv_stethoscope_y"] = conv_stethoscope - normalized_points = normalize_quadrilateral_with_point(source_points.flatten(), stethoscope_point) - row_convert = { - f"{pos}_{coord}": normalized_points[i, j] - for i, pos in enumerate( - ["left_shoulder", "right_shoulder", "left_hip", "right_hip", "stethoscope"] - ) - for j, coord in enumerate(["x", "y"]) - } + if NORMALIZE_ENABLED: + input_data = pd.DataFrame([norm_row]) + else: + input_data = pd.DataFrame([row]) - input_data = pd.DataFrame([row_convert]) - input_columns = list(row_convert.keys()) + input_columns = [ + f"{pos}_{coord}" + for pos in ["left_shoulder", "right_shoulder", "left_hip", "right_hip", "stethoscope"] + for coord in ["x", "y"] + ] if LIGHTGBM_ENABLED: - row["lightGBM_stethoscope_x"] = int(lgb_model_x.predict(input_data[input_columns])[0]) - row["lightGBM_stethoscope_y"] = int(lgb_model_y.predict(input_data[input_columns])[0]) + lgb_x = int(lgb_model_x.predict(input_data[input_columns])[0]) + lgb_y = int(lgb_model_y.predict(input_data[input_columns])[0]) + row["lightGBM_stethoscope_x"], row["lightGBM_stethoscope_y"] = lgb_x, lgb_y + norm_row["lightGBM_stethoscope_x"], norm_row["lightGBM_stethoscope_y"] = lgb_x, lgb_y if XGBOOST_ENABLED: - row["Xgboost_stethoscope_x"] = int(xg_model_x.predict(input_data[input_columns])[0]) - row["Xgboost_stethoscope_y"] = int(xg_model_y.predict(input_data[input_columns])[0]) + xg_x = int(xg_model_x.predict(input_data[input_columns])[0]) + xg_y = int(xg_model_y.predict(input_data[input_columns])[0]) + row["Xgboost_stethoscope_x"], row["Xgboost_stethoscope_y"] = xg_x, xg_y + norm_row["Xgboost_stethoscope_x"], norm_row["Xgboost_stethoscope_y"] = xg_x, xg_y 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) + norm_writer.writerow(norm_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}") + print(f"Processed and saved normalized results to: {normalized_csv_path}") generate_visualizations(csv_path, base_dir, output_dir) else: