Spaces:
Running
Running
| 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). The official home of YOLO-NAS-Pose is SuperGradients, [gives us a ⭐️ on GitHub!](https://github.com/Deci-AI/super-gradients)", | |
| live=True, | |
| allow_flagging=False, | |
| ) | |
| # Launch the interface | |
| iface.launch() |