TwolFace/Head-pose-estimation

1from fastapi import FastAPI, UploadFile, File, HTTPException2from fastapi.responses import FileResponse3import tempfile4import os5from pathlib import Path6from model import predict_pose_video, predict_pose_image7 8app = FastAPI(title="Head Pose Estimation API", version="1.0.0")9 10 11 12@app.get("/")13async def read_root():14    return {15        "message": "Head Pose Estimation API",16        "endpoint": "/predict - Upload image/video for pose estimation (auto-detects type)",17        "supported_formats": {18            "images": ["jpg", "jpeg", "png", "gif", "bmp", "tiff", "ico", "webp"],19            "videos": ["mp4", "avi", "mov", "mkv"]20        }21    }22 23@app.post("/predict")24async def predict(file: UploadFile = File(...)):25    """26    Unified endpoint for both image and video processing.27    Automatically detects file type and processes accordingly.28    """29    # Get file extension and content30    file_content = await file.read()31    file_extension = file.filename.split('.')[-1].lower()32    33    # Define supported file types34    video_extensions = {'mp4', 'avi', 'mov', 'mkv'}35    photo_extensions = {'jpg', 'jpeg', 'png', 'gif', 'bmp', 'tiff', 'ico', 'webp'}36    37    # Validate file type38    if file_extension not in video_extensions and file_extension not in photo_extensions:39        raise HTTPException(40            status_code=400,41            detail=f"File type .{file_extension} not allowed. Allowed types are: {', '.join(video_extensions | photo_extensions)}"42        )43    44    try:45        if file_extension in video_extensions:46            # Create temporary file for output47            with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_output:48                output_path = temp_output.name49            50            # Process the video51            result = predict_pose_video(file_content, output_path)52            53            if not result["success"]:54                if os.path.exists(output_path):55                    os.unlink(output_path)56                raise HTTPException(57                    status_code=500,58                    detail=result.get("error", "Unknown error during video processing")59                )60            61            # Return the processed video file62            return FileResponse(63                path=output_path,64                media_type='video/mp4',65                filename=f"processed_{file.filename}",66                headers={67                    "X-Processing-Stats": (68                        f"Frames: {result['total_frames']}, "69                        f"Faces: {result['faces_detected']}, "70                        f"Detection Rate: {result['face_detection_rate']:.1f}%"71                    )72                },73                background=lambda: os.unlink(output_path) if os.path.exists(output_path) else None74            )75            76        else:  # photo_extensions77            # Process image78            with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_extension}") as temp_file:79                temp_file.write(file_content)80                temp_path = temp_file.name81            82            try:83                result = predict_pose_image(temp_path)84                85                if not result["success"]:86                    raise HTTPException(87                        status_code=500,88                        detail=result.get("error", "Failed to process image")89                    )90                91                # Return the processed image file92                return FileResponse(93                    path=result["output_path"],94                    media_type=f"image/{file_extension}",95                    filename=f"processed_{file.filename}",96                    headers={97                        "X-Processing-Stats": (98                            f"Angles -> Pitch: {result['angles']['pitch']:.1f}°, "99                            f"Yaw: {result['angles']['yaw']:.1f}°, "100                            f"Roll: {result['angles']['roll']:.1f}°"101                        )102                    },103                    background=lambda: os.unlink(result["output_path"]) if os.path.exists(result["output_path"]) else None104                )105                106            finally:107                if os.path.exists(temp_path):108                    os.unlink(temp_path)109                    110    except HTTPException:111        raise112    except Exception as e:113        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")114 115 116 117