Kiki05/open_pose

1#!/usr/bin/env python2 3from _future_ import annotations4 5import os6import pathlib7import shlex8import subprocess9import sys10import cv211import gradio as gr12import huggingface_hub13import numpy as np14import openpose as op15 16TITLE = 'Final Year Project - Human Pose Estimation'17DESCRIPTION = 'Pose Estimation Interface | Please provide an input image to estimate & process the pose'18 19HF_TOKEN = os.getenv('HF_TOKEN')20 21 22def load_sample_images() -> list[pathlib.Path]:23    image_dir = pathlib.Path('images')24    if not image_dir.exists():25        image_dir.mkdir()26        dataset_repo = 'hysts/input-images'27        filenames = ['002.tar']28        for name in filenames:29            path = huggingface_hub.hf_hub_download(dataset_repo,30                                                   name,31                                                   repo_type='dataset',32                                                   use_auth_token=HF_TOKEN)33            with tarfile.open(path) as f:34                f.extractall(image_dir.as_posix())35    return sorted(image_dir.rglob('*.jpg'))36 37 38def run(image: np.ndarray, model_complexity: str, enable_segmentation: bool,39        min_detection_confidence: float, background_color: str) -> np.ndarray:40 41    model_path = "path/to/openpose/models"42    params = {"model_folder": model_path, "model_pose": model_complexity}43    opWrapper = op.WrapperPython()44    opWrapper.configure(params)45    opWrapper.start()46 47    datum = op.Datum()48    datum.cvInputData = image49    opWrapper.emplaceAndPop([datum])50 51    res = datum.cvOutputData[:, :, ::-1].copy()52 53    if enable_segmentation:54        if background_color == 'white':55            bg_color = 25556        elif background_color == 'black':57            bg_color = 058        elif background_color == 'green':59            bg_color = (0, 255, 0)  # type: ignore60        else:61            raise ValueError62 63        if datum.cvOutputDataMask is not None:64            res[datum.cvOutputDataMask <= 0.1] = bg_color65        else:66            res[:] = bg_color67 68    return res[:, :, ::-1]69 70 71model_complexities = ['BODY_25', 'COCO', 'MPI']72background_colors = ['white', 'black', 'green']73 74image_paths = load_sample_images()75examples = [[76    path.as_posix(), model_complexities[1], True, 0.5, background_colors[0]77] for path in image_paths]78 79gr.Interface(80    fn=run,81    inputs=[82        gr.Image(label='Input', type='numpy'),83        gr.Radio(label='Model Complexity',84                 choices=model_complexities,85                 type='value',86                 value=model_complexities[1]),87        gr.Checkbox(default=True, label='Enable Segmentation'),88        gr.Slider(label='Minimum Detection Confidence',89                  minimum=0,90                  maximum=1,91                  step=0.05,92                  value=0.5),93        gr.Radio(label='Background Color',94                 choices=background_colors,95                 type='value',96                 value=background_colors[0]),97    ],98    outputs=gr.Image(label='Output', type='numpy'),99    title=TITLE,100    description=DESCRIPTION,101).launch(show_api=False)102