fixing

app.py

@@ -1,99 +1,103 @@
 import cv2
 import streamlit as st
 import tempfile
-import time
 import numpy as np
 from face_detection import FaceDetector
 from mark_detection import MarkDetector
 from pose_estimation import PoseEstimator
 from utils import refine
 
+st.title("Pose Estimation")
 
-st.title("Pose-estimation")
-
+# Choose between Image or Video file upload
 file_type = st.selectbox("Choose the type of file you want to upload", ("Image", "Video"))
-if file_type == "Image":
-    uploaded_file = st.file_uploader("Upload an image of your face", type=["jpg","jpeg", "png"])
-else:
-    uploaded_video = st.file_uploader("Upload a video of your face", type=["mp4","mov","avi","mkv"])
+uploaded_file = st.file_uploader(
+    "Upload an image or video file of your face", 
+    type=["jpg", "jpeg", "png", "mp4", "mov", "avi", "mkv"]
+)
+
+# Display placeholder for real-time video output
+FRAME_WINDOW = st.image([])
 
 if uploaded_file is not None:
+    # Video processing
     if file_type == "Video":
         tfile = tempfile.NamedTemporaryFile(delete=False)
         tfile.write(uploaded_file.read())
         cap = cv2.VideoCapture(tfile.name)
-        print(f"Video source: {tfile.name}")
+        st.write(f"Video source: {tfile.name}")
 
-        #getting frame sizes
+        # Getting frame sizes
         frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
         
-        #face detection
+        # Initialize face detection, landmark detection, and pose estimation models
         face_detector = FaceDetector("assets/face_detector.onnx")
-        #landmark detection
         mark_detector = MarkDetector("assets/face_landmarks.onnx")
-        #pose estimation 
         pose_estimator = PoseEstimator(frame_width, frame_height)
 
-        tm = cv2.TickMeter()
-
-        while True:
-
-            # Read a frame.
-            frame_got, frame = cap.read()
-            if frame_got is False:
+        # Process each frame
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
                 break
 
-
-            # Step 1: Get faces from current frame.
+            # Step 1: Detect faces in the frame
             faces, _ = face_detector.detect(frame, 0.7)
 
-            # Any valid face found?
+            # If a face is detected, proceed with pose estimation
             if len(faces) > 0:
-                tm.start()
-
-                # Step 2: Detect landmarks. Crop and feed the face area into the
-                # mark detector. Note only the first face will be used for
-                # demonstration.
+                # Detect landmarks for the first face
                 face = refine(faces, frame_width, frame_height, 0.15)[0]
                 x1, y1, x2, y2 = face[:4].astype(int)
                 patch = frame[y1:y2, x1:x2]
 
-                # Run the mark detection.
+                # Run landmark detection and convert local face area to global image
                 marks = mark_detector.detect([patch])[0].reshape([68, 2])
-
-                # Convert the locations from local face area to the global image.
                 marks *= (x2 - x1)
                 marks[:, 0] += x1
                 marks[:, 1] += y1
 
-                # Step 3: Try pose estimation with 68 points.
+                # Pose estimation with the detected landmarks
                 pose = pose_estimator.solve(marks)
 
-                tm.stop()
-
-                # All done. The best way to show the result would be drawing the
-                # pose on the frame in realtime.
-
-                # Do you want to see the pose annotation?
+                # Draw the pose on the frame
                 pose_estimator.visualize(frame, pose, color=(0, 255, 0))
 
-                # Do you want to see the axes?
-                # pose_estimator.draw_axes(frame, pose)
+            # Convert frame to RGB for Streamlit display
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            FRAME_WINDOW.image(frame_rgb)
 
-                # Do you want to see the marks?
-                # mark_detector.visualize(frame, marks, color=(0, 255, 0))
+        cap.release()
 
-                # Do you want to see the face bounding boxes?
-                # face_detector.visualize(frame, faces)
+    # Image processing
+    elif file_type == "Image":
+        # Load and process uploaded image
+        image = np.array(Image.open(uploaded_file))
+        frame_height, frame_width, _ = image.shape
 
-            # Draw the FPS on screen.
-            cv2.rectangle(frame, (0, 0), (90, 30), (0, 0, 0), cv2.FILLED)
-            cv2.putText(frame, f"FPS: {tm.getFPS():.0f}", (10, 20),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255))
-
-            # Show preview.
-            cv2.imshow("Preview", frame)
-            if cv2.waitKey(1) == 27:
-                break
+        # Initialize models for detection and pose estimation
+        face_detector = FaceDetector("assets/face_detector.onnx")
+        mark_detector = MarkDetector("assets/face_landmarks.onnx")
+        pose_estimator = PoseEstimator(frame_width, frame_height)
 
+        # Detect face and landmarks
+        faces, _ = face_detector.detect(image, 0.7)
+        if len(faces) > 0:
+            face = refine(faces, frame_width, frame_height, 0.15)[0]
+            x1, y1, x2, y2 = face[:4].astype(int)
+            patch = image[y1:y2, x1:x2]
+
+            # Detect landmarks and map them to global image coordinates
+            marks = mark_detector.detect([patch])[0].reshape([68, 2])
+            marks *= (x2 - x1)
+            marks[:, 0] += x1
+            marks[:, 1] += y1
+
+            # Estimate pose and visualize on image
+            pose = pose_estimator.solve(marks)
+            pose_estimator.visualize(image, pose, color=(0, 255, 0))
+
+            # Convert image to RGB and display in Streamlit
+            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            st.image(image_rgb, caption="Pose Estimated Image", use_column_width=True)

face_detection.py

@@ -36,8 +36,8 @@ class FaceDetector:
         assert os.path.exists(model_file), f"File not found: {model_file}"
         self.center_cache = {}
         self.nms_threshold = 0.4
-        self.session = onnxruntime.InferenceSession(
-            model_file, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+        self.session = onnxruntime.InferenceSession(model_file, providers=['CPUExecutionProvider'])
+
 
         # Get model configurations from the model file.
         # What is the input like?

mark_detection.py

@@ -17,7 +17,7 @@ class MarkDetector:
         assert os.path.exists(model_file), f"File not found: {model_file}"
         self._input_size = 128
         self.model = ort.InferenceSession(
-            model_file, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+            model_file, providers=["CPUExecutionProvider"])
 
     def _preprocess(self, bgrs):
         """Preprocess the inputs to meet the model's needs.

requirements.txt

@@ -1,6 +1,5 @@
 opencv-python-headless
 numpy
-tempfile
 time 
 onnxruntime
 os