diff --git a/HPEModule.py b/HPEModule.py index 875e9db..5f78026 100644 --- a/HPEModule.py +++ b/HPEModule.py @@ -1,3 +1,5 @@ +import os + import cv2 import mediapipe as mp from DataStreamModule import DataStreamModule @@ -13,10 +15,25 @@ class HPEModule(IHPEModule): # This method starts HPE using a camera specified by its name def startHPEwithCamera(self, camera_name): + out = None + # check if the camera_name is a file path or not + if os.path.isfile(camera_name): + # open the video file using cv2.VideoCapture + cap = cv2.VideoCapture(camera_name) + # set the output video file path + output_path = os.path.splitext(camera_name)[0] + "_output.mp4" + # get the frame rate and size of the input video + fps = cap.get(cv2.CAP_PROP_FPS) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - # initialize data stream and camera object - data_stream = DataStreamModule() - cap = data_stream.get_camera_stream(camera_name) + # initialize video writer to save the output video + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) + else: + # initialize data stream and camera object + data_stream = DataStreamModule() + cap = data_stream.get_camera_stream(camera_name) # set the window name using the camera name window_name = f"Pose Estimation on Camera {camera_name}" @@ -25,7 +42,14 @@ class HPEModule(IHPEModule): with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose: while True: # get the next frame from the camera - frame = next(cap) + if out is not None: + ret, frame = cap.read() + if not ret: + break + else: + frame = next(cap) + + # Recolor image to RGB image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) @@ -52,6 +76,10 @@ class HPEModule(IHPEModule): ) cv2.imshow(window_name, image) + + if out is not None: + out.write(image) + keyCode = cv2.waitKey(1) if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1: break diff --git a/PoseEstimationGUI.py b/PoseEstimationGUI.py index 1dbb0f4..481185e 100644 --- a/PoseEstimationGUI.py +++ b/PoseEstimationGUI.py @@ -1,4 +1,6 @@ import tkinter as tk +from tkinter import filedialog +from tkinter import Label from DataStreamModule import DataStreamModule from HPEModule import HPEModule import threading @@ -7,8 +9,10 @@ import threading class PoseEstimationGUI: def __init__(self): + self.mp4_file_path = None self.root = tk.Tk() self.root.title("Pose Estimation GUI") + self.root.geometry("400x200") self.camera_options = [] # Initialize camera options list @@ -23,15 +27,41 @@ class PoseEstimationGUI: self.start_button = tk.Button(self.root, text="Start Pose Estimation", command=self.start_pose_estimation) self.start_button.pack() + # Create a button to select an mp4 file + self.select_file_button = tk.Button(self.root, text="Select MP4 File to Start HPE on", command=self.select_file) + self.select_file_button.pack() + + # Create label widget to indicate file selection status + self.file_selected_label = Label(self.root, text="No file selected", fg="red") + self.file_selected_label.pack(pady=10) + self.root.mainloop() def start_pose_estimation(self): - # Start pose estimation on the selected camera - # print(self.selected_camera.get()) - pose_estimator = HPEModule() - camera = int(self.selected_camera.get()) - pose_thread = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(camera,)) - pose_thread.start() + if self.mp4_file_path is not None: + # Start pose estimation on the selected mp4 file + pose_estimator = HPEModule() + pose_thread_file = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(self.mp4_file_path,)) + self.mp4_file_path = None + self.file_selected_label.config(text="No File selected", fg="red") + pose_thread_file.start() + else: + # Start pose estimation on the selected camera + # print(self.selected_camera.get()) + pose_estimator = HPEModule() + camera = int(self.selected_camera.get()) + pose_thread_camera = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(camera,)) + pose_thread_camera.start() + + def select_file(self): + self.mp4_file_path = filedialog.askopenfilename(initialdir="/", title="Select MP4 File to Start HPE on", + filetypes=(("mp4 files", "*.mp4"), ("all files", "*.*"))) + # Update file selection status label + if self.mp4_file_path: + self.file_selected_label.config(text="File selected", fg="green") + else: + self.file_selected_label.config(text="No File selected", fg="red") + self.mp4_file_path = None if __name__ == '__main__':