Browse Source

Added mp4 functionality

Dev
Erik Hildebrandt 3 years ago
parent
commit
b3572fb730
  1. 36
      HPEModule.py
  2. 42
      PoseEstimationGUI.py

36
HPEModule.py

@ -1,3 +1,5 @@
import os
import cv2 import cv2
import mediapipe as mp import mediapipe as mp
from DataStreamModule import DataStreamModule from DataStreamModule import DataStreamModule
@ -13,10 +15,25 @@ class HPEModule(IHPEModule):
# This method starts HPE using a camera specified by its name # This method starts HPE using a camera specified by its name
def startHPEwithCamera(self, camera_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 # initialize video writer to save the output video
data_stream = DataStreamModule() fourcc = cv2.VideoWriter_fourcc(*'mp4v')
cap = data_stream.get_camera_stream(camera_name) 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 # set the window name using the camera name
window_name = f"Pose Estimation on Camera {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: with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
while True: while True:
# get the next frame from the camera # 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 # Recolor image to RGB
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
@ -52,6 +76,10 @@ class HPEModule(IHPEModule):
) )
cv2.imshow(window_name, image) cv2.imshow(window_name, image)
if out is not None:
out.write(image)
keyCode = cv2.waitKey(1) keyCode = cv2.waitKey(1)
if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1: if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1:
break break

42
PoseEstimationGUI.py

@ -1,4 +1,6 @@
import tkinter as tk import tkinter as tk
from tkinter import filedialog
from tkinter import Label
from DataStreamModule import DataStreamModule from DataStreamModule import DataStreamModule
from HPEModule import HPEModule from HPEModule import HPEModule
import threading import threading
@ -7,8 +9,10 @@ import threading
class PoseEstimationGUI: class PoseEstimationGUI:
def __init__(self): def __init__(self):
self.mp4_file_path = None
self.root = tk.Tk() self.root = tk.Tk()
self.root.title("Pose Estimation GUI") self.root.title("Pose Estimation GUI")
self.root.geometry("400x200")
self.camera_options = [] # Initialize camera options list 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 = tk.Button(self.root, text="Start Pose Estimation", command=self.start_pose_estimation)
self.start_button.pack() 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() self.root.mainloop()
def start_pose_estimation(self): def start_pose_estimation(self):
# Start pose estimation on the selected camera if self.mp4_file_path is not None:
# print(self.selected_camera.get()) # Start pose estimation on the selected mp4 file
pose_estimator = HPEModule() pose_estimator = HPEModule()
camera = int(self.selected_camera.get()) pose_thread_file = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(self.mp4_file_path,))
pose_thread = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(camera,)) self.mp4_file_path = None
pose_thread.start() 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__': if __name__ == '__main__':

Loading…
Cancel
Save