Spaces:
Sleeping
Sleeping
| import cv2 | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| import numpy as np | |
| import math | |
| import tempfile | |
| import os | |
| # Load the model | |
| # Load the model | |
| # Using custom model provided by user | |
| model = YOLO('yolo26x-pose.pt') | |
| def get_angle(p1, p2, p3): | |
| """ | |
| Calculate the angle between three points (p1-p2-p3) at joint p2. | |
| """ | |
| rad = math.atan2(p3[1]-p2[1], p3[0]-p2[0]) - math.atan2(p1[1]-p2[1], p1[0]-p2[0]) | |
| deg = abs(rad * 180.0 / math.pi) | |
| return 360 - deg if deg > 180 else deg | |
| def draw_pose(img, kps): | |
| """ | |
| Draw pose landmarks and classify posture based on knee angle. | |
| """ | |
| legs = [(11, 13, 15), (12, 14, 16)] # hip-knee-ankle indices | |
| status = "Unknown" | |
| for h_idx, k_idx, a_idx in legs: | |
| hip, knee, ankle = kps[h_idx], kps[k_idx], kps[a_idx] | |
| # Check confidence scores (index 2) | |
| if hip[2] > 0.5 and knee[2] > 0.5 and ankle[2] > 0.5: | |
| ang = get_angle(hip[:2], knee[:2], ankle[:2]) | |
| if ang > 160: | |
| posture, color = "STANDING", (0, 255, 0) | |
| elif ang < 140: | |
| posture, color = "SITTING", (255, 0, 0) | |
| else: | |
| posture, color = "BENDING", (0, 165, 255) | |
| # Draw lines and text | |
| pt = lambda p: (int(p[0]), int(p[1])) | |
| cv2.line(img, pt(hip), pt(knee), (255,0,255), 2) | |
| cv2.line(img, pt(knee), pt(ankle), (255,0,255), 2) | |
| label = f"{posture} {int(ang)}" | |
| (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2) | |
| cv2.rectangle(img, (int(hip[0]), int(hip[1]) - 30), (int(hip[0]) + w, int(hip[1])), color, -1) | |
| cv2.putText(img, label, (int(hip[0]), int(hip[1]) - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 255), 2) | |
| return posture | |
| return status | |
| def process_img(img): | |
| """ | |
| Process a single image for posture detection. | |
| """ | |
| if img is None: | |
| return None, "No Image" | |
| # Run inference | |
| res = model(img) | |
| out = img.copy() | |
| msg = "No Person Detected" | |
| if res and res[0].keypoints is not None: | |
| for k in res[0].keypoints.data.cpu().numpy(): | |
| msg = f"POSTURE: {draw_pose(out, k)}" | |
| return out, msg | |
| def process_vid(vid_path): | |
| """ | |
| Process a video file for posture detection. | |
| """ | |
| if not vid_path: | |
| return None | |
| cap = cv2.VideoCapture(vid_path) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # Create temp file for output | |
| fd, out_path = tempfile.mkstemp(suffix='.mp4') | |
| os.close(fd) | |
| # Initialize video writer | |
| writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Run inference | |
| res = model(frame, verbose=False) | |
| if res and res[0].keypoints is not None: | |
| for k in res[0].keypoints.data.cpu().numpy(): | |
| draw_pose(frame, k) | |
| writer.write(frame) | |
| cap.release() | |
| writer.release() | |
| return out_path | |
| # Define Gradio Interface | |
| with gr.Blocks(title="Postures") as app: | |
| gr.Markdown("# Posture Detection\nSimple angle-based classification using YOLOv8: Standing (>160), Sitting (<140)") | |
| with gr.Tab("Image"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| img_inp = gr.Image(type="numpy", label="Input Image") | |
| img_btn = gr.Button("Detect Posture") | |
| with gr.Column(): | |
| img_out = gr.Image(label="Result") | |
| img_stat = gr.Textbox(label="Status") | |
| img_btn.click(process_img, inputs=img_inp, outputs=[img_out, img_stat]) | |
| with gr.Tab("Video"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| vid_inp = gr.Video(label="Input Video") | |
| vid_btn = gr.Button("Process Video") | |
| with gr.Column(): | |
| vid_out = gr.Video(label="Processed Video") | |
| vid_btn.click(process_vid, inputs=vid_inp, outputs=vid_out) | |
| if __name__ == "__main__": | |
| app.launch() | |