From 34d48df51d3dff580f60d0cc77045ba126b9e053 Mon Sep 17 00:00:00 2001 From: Erik Hildebrandt Date: Wed, 26 Apr 2023 21:32:22 +0200 Subject: [PATCH] Added Button to Control Export to GUI --- HPEModule.py | 40 ++++++++++++++++++++++------------------ IHPEModule.py | 2 +- PoseEstimationGUI.py | 14 +++++++++++--- 3 files changed, 34 insertions(+), 22 deletions(-) diff --git a/HPEModule.py b/HPEModule.py index d5833ac..261c746 100644 --- a/HPEModule.py +++ b/HPEModule.py @@ -16,7 +16,7 @@ mp_pose = mp.solutions.pose class HPEModule(IHPEModule): # This method starts HPE using a camera specified by its name - def startHPEwithCamera(self, camera_name): + def startHPEwithCamera(self, camera_name, export_landmarks): out = None # check if the camera_name is a file path or not if os.path.isfile(camera_name): @@ -39,15 +39,17 @@ class HPEModule(IHPEModule): # set the window name using the camera name window_name = f"Pose Estimation on Camera {camera_name}" + print(export_landmarks.get()) + if export_landmarks.get() is True: + # create the root element + root = ET.Element("Landmarks") - # create the root element - root = ET.Element("Landmarks") + # loop through each landmark + for i in range(33): + ET.SubElement(root, f'landmark{i}') - # loop through each landmark - for i in range(33): - ET.SubElement(root, f'landmark{i}') + framecounter = 0 - framecounter = 0 # start pose detection with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose: @@ -75,12 +77,13 @@ class HPEModule(IHPEModule): # Extract landmarks try: landmarks = results.pose_landmarks.landmark - for j, landmark in enumerate(landmarks): - frame_element = ET.SubElement(root.find(f'landmark{j}'), f'frame{framecounter}') - frame_element.set("X", str(landmark.x)) - frame_element.set("Y", str(landmark.y)) - frame_element.set("Z", str(landmark.z)) - framecounter += 1 + if export_landmarks.get() is True: + for j, landmark in enumerate(landmarks): + frame_element = ET.SubElement(root.find(f'landmark{j}'), f'frame{framecounter}') + frame_element.set("X", str(landmark.x)) + frame_element.set("Y", str(landmark.y)) + frame_element.set("Z", str(landmark.z)) + framecounter += 1 # print(landmarks) except: pass @@ -100,8 +103,9 @@ class HPEModule(IHPEModule): if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1: break - xml_string = ET.tostring(root, encoding="UTF-8") - parsed_xml = minidom.parseString(xml_string) - pretty_xml_string = parsed_xml.toprettyxml(indent=" ") - with open(f"{camera_name}_output.xml", "w") as xml_file: - xml_file.write(pretty_xml_string) + if export_landmarks.get() is True: + xml_string = ET.tostring(root, encoding="UTF-8") + parsed_xml = minidom.parseString(xml_string) + pretty_xml_string = parsed_xml.toprettyxml(indent=" ") + with open(f"{camera_name}_output.xml", "w") as xml_file: + xml_file.write(pretty_xml_string) diff --git a/IHPEModule.py b/IHPEModule.py index 27612b7..c331278 100644 --- a/IHPEModule.py +++ b/IHPEModule.py @@ -5,5 +5,5 @@ class IHPEModule(ABC): # This method starts HPE using a camera specified by its name @abstractmethod - def startHPEwithCamera(self, camera_name): + def startHPEwithCamera(self, camera_name, export_landmarks): pass diff --git a/PoseEstimationGUI.py b/PoseEstimationGUI.py index ab60cd1..d684403 100644 --- a/PoseEstimationGUI.py +++ b/PoseEstimationGUI.py @@ -12,7 +12,7 @@ class PoseEstimationGUI: self.mp4_file_path = None self.root = tk.Tk() self.root.title("Pose Estimation GUI") - self.root.geometry("400x250") + self.root.geometry("400x400") self.root.config(bg="#F9F9F9") self.camera_options = [] # Initialize camera options list @@ -42,13 +42,21 @@ class PoseEstimationGUI: font=("Arial", 12, "bold"), command=self.start_pose_estimation) self.start_button.pack(pady=(20, 10)) + # Create a checkbox to enable/disable Export of Landmarks + self.export_Landmarks = tk.BooleanVar() + self.export_Landmarks.set(False) + export_Landmarks_checkbox = tk.Checkbutton(self.root, text="Export Landmarks", variable=self.export_Landmarks, + bg="#F9F9F9", + font=("Arial", 12, "bold"), onvalue=True, offvalue=False) + export_Landmarks_checkbox.pack() + self.root.mainloop() def start_pose_estimation(self): 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,)) + pose_thread_file = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(self.mp4_file_path, self.export_Landmarks,)) self.mp4_file_path = None self.file_selected_label.config(text="No File selected", fg="red") pose_thread_file.start() @@ -57,7 +65,7 @@ class PoseEstimationGUI: # 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 = threading.Thread(target=pose_estimator.startHPEwithCamera, args=(camera, self.export_Landmarks,)) pose_thread_camera.start() def select_file(self):