Atualli/mediapipe-pose-estimation

1#!/usr/bin/env python2 3from __future__ import annotations4 5import pathlib6 7import gradio as gr8import mediapipe as mp9import numpy as np10 11mp_drawing = mp.solutions.drawing_utils12mp_drawing_styles = mp.solutions.drawing_styles13mp_pose = mp.solutions.pose14 15TITLE = 'MediaPipe Human Pose Estimation'16DESCRIPTION = 'https://google.github.io/mediapipe/'17 18 19def run(image: np.ndarray, model_complexity: int, enable_segmentation: bool,20        min_detection_confidence: float, background_color: str) -> np.ndarray:21    with mp_pose.Pose(22            static_image_mode=True,23            model_complexity=model_complexity,24            enable_segmentation=enable_segmentation,25            min_detection_confidence=min_detection_confidence) as pose:26        results = pose.process(image)27 28    res = image[:, :, ::-1].copy()29    if enable_segmentation:30        if background_color == 'white':31            bg_color = 25532        elif background_color == 'black':33            bg_color = 034        elif background_color == 'green':35            bg_color = (0, 255, 0)  # type: ignore36        else:37            raise ValueError38 39        if results.segmentation_mask is not None:40            res[results.segmentation_mask <= 0.1] = bg_color41        else:42            res[:] = bg_color43 44    mp_drawing.draw_landmarks(res,45                              results.pose_landmarks,46                              mp_pose.POSE_CONNECTIONS,47                              landmark_drawing_spec=mp_drawing_styles.48                              get_default_pose_landmarks_style())49 50    return res[:, :, ::-1]51 52 53model_complexities = list(range(3))54background_colors = ['white', 'black', 'green']55 56image_paths = sorted(pathlib.Path('images').rglob('*.jpg'))57examples = [[path, model_complexities[1], True, 0.5, background_colors[0]]58            for path in image_paths]59 60gr.Interface(61    fn=run,62    inputs=[63        gr.Image(label='Input', type='numpy'),64        gr.Radio(label='Model Complexity',65                 choices=model_complexities,66                 type='index',67                 value=model_complexities[1]),68        gr.Checkbox(label='Enable Segmentation', value=True),69        gr.Slider(label='Minimum Detection Confidence',70                  minimum=0,71                  maximum=1,72                  step=0.05,73                  value=0.5),74        gr.Radio(label='Background Color',75                 choices=background_colors,76                 type='value',77                 value=background_colors[0]),78    ],79    outputs=gr.Image(label='Output', height=500),80    examples=examples,81    title=TITLE,82    description=DESCRIPTION,83).queue().launch()84