diff --git a/main.py b/main.py index 3ea2c6a..71369b2 100644 --- a/main.py +++ b/main.py @@ -13,6 +13,15 @@ from util.ears_ai import EarsAI from util.calc_ste_position import CalcStethoscopePosition from modules.EARSForDL.model import RegressionResNet +from dotenv import load_dotenv + +# Load environment variables +load_dotenv() + +# Get colors from environment variables +CONV_COLOR = tuple(map(int, os.getenv("CONV_COLOR", "0,255,0").split(","))) # Default: Green +XGBOOST_COLOR = tuple(map(int, os.getenv("XGBOOST_COLOR", "255,0,0").split(","))) # Default: Red +LIGHTGBM_COLOR = tuple(map(int, os.getenv("LIGHTGBM_COLOR", "0,0,255").split(","))) # Default: Blue def load_model(model_path, model_type="lgb"): @@ -72,7 +81,7 @@ return transform(Image.open(image_path).convert("RGB")).unsqueeze(0) -def process_images(base_dir, draw_trajectory=True): +def process_images(base_dir): ears_ai = EarsAI() calc_position = CalcStethoscopePosition() results_dir = os.path.join(os.path.dirname(base_dir), "results") @@ -184,22 +193,27 @@ writer.writerow(row) print(f"Processed and saved results to: {csv_path}") - generate_visualizations(csv_path, base_dir, draw_trajectory) + generate_visualizations(csv_path, base_dir) else: print("No data to write to CSV.") -def generate_visualizations(csv_path, original_images_dir, draw_trajectory): +def generate_visualizations(csv_path, original_images_dir): df = pd.read_csv(csv_path) body_image = cv2.imread("images/body/BodyF.png") results_dir = "images/body/results" os.makedirs(results_dir, exist_ok=True) dirs = {"marked": "marked_images", "conv": "conv", "Xgboost": "Xgboost", "lightGBM": "lightGBM"} + + # Create all necessary directories for dir_name in dirs.values(): os.makedirs(os.path.join(results_dir, dir_name), exist_ok=True) + os.makedirs(os.path.join(results_dir, f"{dir_name}_with_trajectory"), exist_ok=True) + os.makedirs(os.path.join(results_dir, f"{dir_name}_without_trajectory"), exist_ok=True) points = {key: [] for key in ["conv", "Xgboost", "lightGBM"]} + colors = {"conv": CONV_COLOR, "Xgboost": XGBOOST_COLOR, "lightGBM": LIGHTGBM_COLOR} for _, row in df.iterrows(): original_image = cv2.imread(os.path.join(original_images_dir, row["image_file_name"])) @@ -216,26 +230,41 @@ for key in points: x, y = int(row[f"{key}_stethoscope_x"]), int(row[f"{key}_stethoscope_y"]) points[key].append((x, y)) - image = body_image.copy() - if draw_trajectory and len(points[key]) > 1: - cv2.polylines(image, [np.array(points[key])], False, (0, 255, 0), 2) - cv2.circle(image, (x, y), 10, (0, 255, 0), -1) - cv2.imwrite(os.path.join(results_dir, dirs[key], row["image_file_name"]), image) + + # Create image with trajectory + image_with_trajectory = body_image.copy() + if len(points[key]) > 1: + cv2.polylines(image_with_trajectory, [np.array(points[key])], False, colors[key], 2) + cv2.circle(image_with_trajectory, (x, y), 10, colors[key], -1) + cv2.imwrite( + os.path.join(results_dir, f"{dirs[key]}_with_trajectory", row["image_file_name"]), image_with_trajectory + ) + + # Create image without trajectory + image_without_trajectory = body_image.copy() + cv2.circle(image_without_trajectory, (x, y), 10, colors[key], -1) + cv2.imwrite( + os.path.join(results_dir, f"{dirs[key]}_without_trajectory", row["image_file_name"]), + image_without_trajectory, + ) # Generate videos for key in dirs: - create_video_from_images( - os.path.join(results_dir, dirs[key]), os.path.join(results_dir, f"{key}_video_with_trajectory.mp4") - ) if key != "marked": create_video_from_images( - os.path.join(results_dir, dirs[key]), + os.path.join(results_dir, f"{dirs[key]}_with_trajectory"), + os.path.join(results_dir, f"{key}_video_with_trajectory.mp4"), + ) + create_video_from_images( + os.path.join(results_dir, f"{dirs[key]}_without_trajectory"), os.path.join(results_dir, f"{key}_video_without_trajectory.mp4"), - with_trajectory=False, ) + # Generate video for marked images + create_video_from_images(os.path.join(results_dir, dirs["marked"]), os.path.join(results_dir, "marked_video.mp4")) -def create_video_from_images(image_dir, output_path, with_trajectory=True): + +def create_video_from_images(image_dir, output_path): images = sorted( [img for img in os.listdir(image_dir) if img.endswith(".png")], key=lambda x: int(re.search(r"(\d+)", x).group()), @@ -250,14 +279,8 @@ video = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), 30, (width, height)) - background = cv2.imread("images/body/BodyF.png") if not with_trajectory else None - for image in images: img = cv2.imread(os.path.join(image_dir, image)) - if not with_trajectory: - mask = cv2.threshold(cv2.cvtColor(img, cv2.COLOR_BGR2GRAY), 10, 255, cv2.THRESH_BINARY)[1] - img = cv2.bitwise_and(img, img, mask=mask) - img = cv2.add(img, cv2.bitwise_and(background, background, mask=cv2.bitwise_not(mask))) video.write(img) video.release() @@ -268,9 +291,8 @@ 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="output", 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() video_to_frames(args.video_path, args.output_dir) - process_images(args.output_dir, args.draw_trajectory) + process_images(args.output_dir)