Spaces:
Sleeping
Sleeping
File size: 4,296 Bytes
5751a55 aa54535 5751a55 | 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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | 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()
|