Spaces:
Build error
Build error
File size: 3,342 Bytes
d13f105 99a0703 d13f105 2cdc5a5 202d2bb d13f105 2cdc5a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
from io import BytesIO
import cv2
import gradio as gr
import numpy as np
import requests
from PIL import Image
from super_gradients.common.object_names import Models
from super_gradients.training import models
from super_gradients.training.utils.visualization.detection import draw_bbox
from super_gradients.training.utils.visualization.pose_estimation import PoseVisualization
# Initialize your pose estimation model
yolo_nas_pose = models.get("yolo_nas_pose_l",
num_classes=17,
checkpoint_path="./yolo_nas_pose_l_coco_pose.pth")
def process_and_predict(url=None,
image=None,
confidence=0.5,
iou=0.5):
# If a URL is provided, use it directly for prediction
if url is not None and url.strip() != "":
response = requests.get(url)
image = Image.open(BytesIO(response.content))
image = np.array(image)
result = yolo_nas_pose.predict(image, conf=confidence,iou=iou)
# If a file is uploaded, read it, convert it to a numpy array and use it for prediction
elif image is not None:
result = yolo_nas_pose.predict(image, conf=confidence,iou=iou)
else:
return None # If no input is provided, return None
# Extract prediction data
image_prediction = result._images_prediction_lst[0]
pose_data = image_prediction.prediction
# Visualize the prediction
output_image = PoseVisualization.draw_poses(
image=image_prediction.image,
poses=pose_data.poses,
boxes=pose_data.bboxes_xyxy,
scores=pose_data.scores,
is_crowd=None,
edge_links=pose_data.edge_links,
edge_colors=pose_data.edge_colors,
keypoint_colors=pose_data.keypoint_colors,
joint_thickness=2,
box_thickness=2,
keypoint_radius=5
)
blank_image = np.zeros_like(image_prediction.image)
skeleton_image = PoseVisualization.draw_poses(
image=blank_image,
poses=pose_data.poses,
boxes=pose_data.bboxes_xyxy,
scores=pose_data.scores,
is_crowd=None,
edge_links=pose_data.edge_links,
edge_colors=pose_data.edge_colors,
keypoint_colors=pose_data.keypoint_colors,
joint_thickness=2,
box_thickness=2,
keypoint_radius=5
)
return output_image, skeleton_image
# Define the Gradio interface
iface = gr.Interface(
fn=process_and_predict,
inputs=[
gr.Textbox(placeholder="Enter Image URL", label="Image URL"),
gr.Image(label="Upload Image", type='numpy'),
gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label="Confidence Threshold"),
gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label="IoU Threshold")
],
outputs=[
gr.components.Image(label="Estimated Pose"),
gr.components.Image(label="Skeleton Only")
],
title="YOLO-NAS-Pose Demo",
description="Upload an image, enter an image URL, or use your webcam to use a pretrained YOLO-NAS-Pose L for inference. Get more hands-on with the [starter notebook for inference](https://bit.ly/yn-pose-inference), and learn how to fine-tune your own model with the [fine-tuning notebook](https://bit.ly/yn-pose-fine-tuning).",
live=True,
allow_flagging=False,
)
# Launch the interface
iface.launch() |