preller/test-pose

1import os2import tempfile3import uuid4import shutil5from pathlib import Path6import cv27import mediapipe as mp8import numpy as np9import time10import gc11from fastapi import FastAPI, File, UploadFile12from fastapi.responses import FileResponse13import uvicorn14 15app = FastAPI(title="MediaPipe Pose Detection API")16 17# Initialize MediaPipe Pose18mp_pose = mp.solutions.pose19mp_drawing = mp.solutions.drawing_utils20mp_drawing_styles = mp.solutions.drawing_styles21 22# Set up temporary directory for processing23TEMP_DIR = Path("./temp_files")24TEMP_DIR.mkdir(exist_ok=True)25 26@app.post("/process-video/")27async def process_video(file: UploadFile = File(...)):28    """29    Process a video with MediaPipe pose detection.30    Upload a video file and receive a processed video with pose landmarks.31    """32    # Create a unique filename33    temp_input_file = TEMP_DIR / f"input_{uuid.uuid4()}.mp4"34    temp_output_file = TEMP_DIR / f"output_{uuid.uuid4()}.mp4"35    temp_frames_dir = TEMP_DIR / f"frames_{uuid.uuid4()}"36    temp_frames_dir.mkdir(exist_ok=True)37    38    try:39        # Save uploaded file40        with open(temp_input_file, "wb") as f:41            shutil.copyfileobj(file.file, f)42        43        # Process the video44        success = process_video_file(45            input_path=str(temp_input_file),46            output_path=str(temp_output_file),47            frames_dir=str(temp_frames_dir)48        )49        50        if not success:51            return {"error": "Failed to process video"}52        53        # Return the processed video54        return FileResponse(55            path=temp_output_file,56            media_type="video/mp4",57            filename="processed_video.mp4"58        )59    60    except Exception as e:61        return {"error": str(e)}62    63    finally:64        # Clean up temporary files65        if temp_input_file.exists():66            temp_input_file.unlink()67        if temp_frames_dir.exists():68            shutil.rmtree(temp_frames_dir)69        # Keep output file until it's served to the client70 71def process_video_file(input_path, output_path, frames_dir):72    """73    Process a video file using MediaPipe pose detection.74    This is adapted from 2posetest-fixed.py.75    """76    # Extract frames77    success, width, height, fps, frame_count = extract_frames(input_path, frames_dir)78    if not success:79        return False80    81    # Process frames82    success = process_frames(width, height, fps, frame_count, frames_dir, output_path)83    84    # Clean up85    cleanup_temp(frames_dir)86    87    return success88 89def extract_frames(input_video_path, temp_dir):90    """Extract frames from video to temporary directory"""91    cap = cv2.VideoCapture(input_video_path)92    if not cap.isOpened():93        return False, None, None, None, 094    95    # Get video properties96    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))97    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))98    fps = cap.get(cv2.CAP_PROP_FPS)99    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))100    101    frame_idx = 0102    while True:103        success, frame = cap.read()104        if not success:105            break106        107        frame_path = os.path.join(temp_dir, f"frame_{frame_idx:04d}.jpg")108        cv2.imwrite(frame_path, frame)109        frame_idx += 1110    111    cap.release()112    return True, width, height, fps, frame_idx113 114def process_frames(width, height, fps, frame_count, temp_dir, output_video_path):115    """Process frames with MediaPipe and create output video"""116    # Configure pose detection with lower resource usage117    pose = mp_pose.Pose(118        static_image_mode=True,  # Using static image mode for better results119        model_complexity=0,      # Lower model complexity120        enable_segmentation=False,121        min_detection_confidence=0.3,122        min_tracking_confidence=0.3123    )124    125    # Create video writer126    fourcc = cv2.VideoWriter_fourcc(*'mp4v')127    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))128    129    if not out.isOpened():130        return False131    132    # Process in smaller batches to manage memory133    batch_size = 10134    processed_frames = 0135    136    for frame_idx in range(frame_count):137        # Clear memory periodically138        if frame_idx % batch_size == 0:139            gc.collect()140            141        frame_path = os.path.join(temp_dir, f"frame_{frame_idx:04d}.jpg")142        if not os.path.exists(frame_path):143            continue144        145        # Read frame from file146        image = cv2.imread(frame_path)147        if image is None:148            continue149            150        # Process with MediaPipe151        try:152            # Reduce image resolution if needed for performance153            if width > 640:154                scale_factor = 640.0 / width155                working_image = cv2.resize(image, (0, 0), fx=scale_factor, fy=scale_factor)156            else:157                working_image = image158                159            # Convert to RGB for MediaPipe160            image_rgb = cv2.cvtColor(working_image, cv2.COLOR_BGR2RGB)161            162            # Process frame163            results = pose.process(image_rgb)164            165            # Create annotated image166            annotated_image = image.copy()167            168            # Draw pose landmarks if detected169            if results.pose_landmarks:170                # If we processed a scaled image, adjust landmarks back to original size171                if width > 640:172                    # Create a temporary image for drawing173                    temp_annotated = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)174                    mp_drawing.draw_landmarks(175                        temp_annotated,176                        results.pose_landmarks,177                        mp_pose.POSE_CONNECTIONS,178                        landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style()179                    )180                    # Resize back to original dimensions181                    annotated_image = cv2.resize(temp_annotated, (width, height))182                else:183                    # Direct drawing on original sized image184                    mp_drawing.draw_landmarks(185                        annotated_image,186                        results.pose_landmarks,187                        mp_pose.POSE_CONNECTIONS,188                        landmark_drawing_spec=mp_drawing_styles.get_default_pose_landmarks_style()189                    )190            191            # Write to video192            out.write(annotated_image)193            processed_frames += 1194            195        except Exception as e:196            print(f"Error processing frame {frame_idx}: {e}")197    198    # Release resources199    out.release()200    pose = None201    gc.collect()202    203    return processed_frames > 0204 205def cleanup_temp(temp_dir):206    """Remove temporary frame files"""207    try:208        for file in os.listdir(temp_dir):209            file_path = os.path.join(temp_dir, file)210            if os.path.isfile(file_path):211                os.remove(file_path)212    except:213        pass214 215@app.get("/")216async def root():217    return {"message": "Upload a video to /process-video/ to apply MediaPipe pose detection"}218 219if __name__ == "__main__":220    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)