diff --git a/HPEModule.py b/HPEModule.py index 33f1a4e..d5833ac 100644 --- a/HPEModule.py +++ b/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 @@ -38,8 +40,18 @@ class HPEModule(IHPEModule): # set the window name using the camera name window_name = f"Pose Estimation on Camera {camera_name}" + # 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: @@ -63,7 +75,13 @@ class HPEModule(IHPEModule): # Extract landmarks try: landmarks = results.pose_landmarks.landmark - #print(landmarks) + 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 @@ -81,3 +99,9 @@ class HPEModule(IHPEModule): keyCode = cv2.waitKey(1) 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)