Added Button to Control Export to GUI

HPEModule.py

@@ -1,9 +1,11 @@
 import os
+from xml.dom import minidom
 
 import cv2
 import mediapipe as mp
 from DataStreamModule import DataStreamModule
 from IHPEModule import IHPEModule
+import xml.etree.ElementTree as ET
 
 # initialize mediapipe drawing utilities and pose models
 mp_drawing = mp.solutions.drawing_utils
@@ -14,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):
@@ -37,9 +39,21 @@ 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")
+
+            # loop through each landmark
+            for i in range(33):
+                ET.SubElement(root, f'landmark{i}')
+
+            framecounter = 0
+
 
         # start pose detection
         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
                 if out is not None:
@@ -49,8 +63,6 @@ class HPEModule(IHPEModule):
                 else:
                     frame = next(cap)
 
-
-
                 # Recolor image to RGB
                 image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                 image.flags.writeable = False
@@ -65,7 +77,14 @@ class HPEModule(IHPEModule):
                 # Extract landmarks
                 try:
                     landmarks = results.pose_landmarks.landmark
-                    print(landmarks)
+                    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
 
@@ -84,5 +103,9 @@ class HPEModule(IHPEModule):
                 if cv2.getWindowProperty(window_name, cv2.WND_PROP_VISIBLE) < 1:
                     break
 
-            # destroy the window after exiting the loop
+        if export_landmarks.get() is True:
-            #cv2.destroyWindow(window_name)
+            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)

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

PoseEstimationGUI.py

@@ -12,7 +12,8 @@ class PoseEstimationGUI:
         self.mp4_file_path = None
         self.root = tk.Tk()
         self.root.title("Pose Estimation GUI")
-        self.root.geometry("400x200")
+        self.root.geometry("400x400")
+        self.root.config(bg="#F9F9F9")
 
         self.camera_options = []  # Initialize camera options list
 
@@ -20,20 +21,34 @@ class PoseEstimationGUI:
         self.camera_options = DataStreamModule().findCameras()
         self.selected_camera = tk.StringVar()
         self.selected_camera.set(self.camera_options[0])  # default value
+        camera_label = tk.Label(self.root, text="Select a camera:", bg="#F9F9F9", font=("Arial", 12, "bold"))
+        camera_label.pack(pady=(20, 0))
         camera_menu = tk.OptionMenu(self.root, self.selected_camera, *self.camera_options)
+        camera_menu.config(font=("Arial", 12))
         camera_menu.pack()
 
-        # Create a button to start pose estimation
-        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 = tk.Button(self.root, text="Select MP4 File to Start HPE on", bg="#2196F3", fg="white",
-        self.select_file_button.pack()
+                                             font=("Arial", 12, "bold"), command=self.select_file)
+        self.select_file_button.pack(pady=(10, 20))
 
         # 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 = Label(self.root, text="No file selected", fg="red", bg="#F9F9F9",
-        self.file_selected_label.pack(pady=10)
+                                          font=("Arial", 12))
+        self.file_selected_label.pack(pady=(10, 0))
+
+        # Create a button to start pose estimation
+        self.start_button = tk.Button(self.root, text="Start Pose Estimation", bg="#4CAF50", fg="white",
+                                      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()
 
@@ -41,7 +56,7 @@ class PoseEstimationGUI:
         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()
@@ -50,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):