import argparse
import os
import sys
import time
import cv2
import matplotlib.pyplot as plt
import numpy as np
import requests
import kivy
from kivy.app import App
from kivy.uix.label import Label
from kivy.uix.button import Button
from kivy.uix.boxlayout import BoxLayout
from kivy.uix.checkbox import CheckBox
yolov7_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'vendor/yolov7')
sys.path.append(yolov7_path)
from models.experimental import attempt_load
import torch
from torchvision import transforms
from utils.datasets import letterbox
from utils.general import non_max_suppression_kpt, strip_optimizer, xyxy2xywh
from utils.plots import colors, output_to_keypoint, plot_one_box_kpt, plot_skeleton_kpts
from utils.torch_utils import select_device
# from tts.tttest import generate_audios, play_audio
import asyncio
import threading
import websockets
import json
HOST = 'localhost'
PORT = '8000'
POSEWEIGHTS = 'src_models/yolov7-w6-pose.pt'
async def predict_request(payload):
"""
Args:
- payload: {'array': (1, 10, 50) shape (10 frames)}
Returns:
- score: Value between 0 and 1
- status: Good or bad posture (depending on threshold:0.7)
"""
uri = f"ws://{HOST}:{PORT}"
try:
async with websockets.connect(uri) as ws:
payload_json = json.dumps(payload)
await ws.send(payload_json)
raw_prediction = await ws.recv()
prediction = json.loads(raw_prediction)
score = prediction['score']
status = prediction['status']
return score, status
except:
return None, 'server-error'
def predict_http_request(payload):
response = requests.post(f"http://{HOST}:{PORT}/predict", json=payload)
if response.status_code == 200:
return response.json()
else:
print("Error:", response.status_code)
print(response.text)
@torch.no_grad()
def run(source, device, separation, length, multiple):
current_score = 0
current_status = 'good'
# global ap_model
separation = int(separation)
length = int(length)
frame_count = 0 #count no of frames
total_fps = 0 #count total fps
device = select_device(opt.device) #select device
model = attempt_load(POSEWEIGHTS, map_location=device) #Load model
_ = model.eval()
names = model.module.names if hasattr(model, 'module') else model.names # get class names
if source.isnumeric() :
cap = cv2.VideoCapture(int(source)) #pass video to videocapture object
else:
cap = cv2.VideoCapture(source) #pass video to videocapture object
if (cap.isOpened() == False): #check if videocapture not opened
print('Error while trying to read video. Please check path again')
raise SystemExit()
else:
frame_width = int(cap.get(3)) #get video frame width
# logic for multiple persons
people = {}
next_object_id = 0
# logic for single persons
current_sequence = []
previous_status = "None"
longevity = 0 # frames spent in the current status
# generate_audios("good"); generate_audios("bad")
# bad_audio_thread = threading.Thread(target=play_audio, args=["bad"])
empty = False
while(cap.isOpened):
ret, frame = cap.read()
if ret:
orig_image = frame
image = cv2.cvtColor(orig_image, cv2.COLOR_BGR2RGB)
image = letterbox(image, (frame_width), stride=64, auto=True)[0]
image = transforms.ToTensor()(image)
image = torch.tensor(np.array([image.numpy()]))
image = image.to(device)
image = image.float()
with torch.no_grad(): #get predictions
output_data, _ = model(image)
output_data = non_max_suppression_kpt(output_data, #Apply non max suppression
0.25, # Conf. Threshold.
0.65, # IoU Threshold.
nc=model.yaml['nc'], # Number of classes.
nkpt=model.yaml['nkpt'], # Number of keypoints.
kpt_label=True)
output = output_to_keypoint(output_data)
if multiple:
if len(output) == 0:
if not empty:
print("Wiping data, waiting for objects to appear in frame")
people = {}
next_object_id = 0
empty = True
else:
empty = False
else:
if output.shape[0] > 0:
if frame_count % separation == 0:
landmarks = output[0, 7:].T
current_sequence += [landmarks[:-1]]
if len(current_sequence) == 10:
current_sequence = np.array([current_sequence])
payload = {'array': current_sequence.tolist() }
response = predict_http_request(payload)
current_score = response['score']
previous_status = current_status
current_status = response['status']
# score, status = asyncio.run(predict_request(payload))
# if status == 'server-error':
# print('Server error or server not launched')
# print(score, status)
current_sequence = []
# if current_status == previous_status:
# if not bad_audio_thread.is_alive() and longevity < 30:
# longevity += 1
# else:
# longevity = 0
# else:
# longevity = 0
# if longevity == 30 and current_status == "bad":
# try:
# if not bad_audio_thread.is_alive():
# bad_audio_thread = threading.Thread(target=play_audio("bad"))
# bad_audio_thread.start()
# except Exception as e:
# pass
im0 = image[0].permute(1, 2, 0) * 255 # Change format [b, c, h, w] to [h, w, c] for displaying the image.
im0 = im0.cpu().numpy().astype(np.uint8)
im0 = cv2.cvtColor(im0, cv2.COLOR_RGB2BGR) #reshape image format to (BGR)
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
for i, pose in enumerate(output_data): # detections per image
if empty: break
if len(output_data) == 0:
continue
for det_index, (*xyxy, conf, cls) in enumerate(reversed(pose[:,:6])): #loop over poses for drawing on frame
c = int(cls) # integer class
kpts = pose[det_index, 6:]
if multiple:
# get the centroid (cx, cy) for the current rectangle
rect = [tensor.cpu().numpy() for tensor in xyxy]
cx, cy = (rect[0] + rect[2]) / 2, (rect[1] + rect[3]) / 2
matched_object_id = None
# iterating through known people
for object_id, data in people.items():
distance = np.sqrt((cx - data['centroid'][0]) ** 2 + (cy - data['centroid'][1]) ** 2)
print(distance)
if distance < 300: # Adjust the threshold as needed
matched_object_id = object_id
break
if matched_object_id is None:
matched_object_id = next_object_id
next_object_id += 1
if matched_object_id not in people:
people[matched_object_id] = {'centroid': (cx, cy), 'yoloid': det_index, 'status': 'good', 'score': 0, 'sequence' : []}
else:
people[matched_object_id]['centroid'] = (cx, cy)
people[matched_object_id]['yoloid'] = det_index
obj = people[matched_object_id]
label = f"ID #{obj['yoloid']} Score: {obj['score']:.2f}"
plot_one_box_kpt(xyxy, im0, label=label, color=colors(c, True),
line_thickness=3, kpt_label=True, kpts=kpts, steps=3,
cmap=people[matched_object_id]['status'])
else:
label = f"ID #{0} Score: {current_score:.2f}"
plot_one_box_kpt(xyxy, im0, label=label, color=colors(c, True),
line_thickness=3,kpt_label=True, kpts=kpts, steps=3,
cmap=current_status)
if frame_count % separation == 0 and multiple:
for _, data in people.items():
if data['yoloid'] < output.shape[0]:
yoloid = data['yoloid']
landmarks = output[yoloid, 7:].T
data['sequence'] += [landmarks[:-1]]
if len(data['sequence']) == length:
payload = {'array': np.array([data['sequence']]).tolist()}
response = predict_http_request(payload)
data['score'] = response['score']
data['status'] = response['status']
data['sequence'] = []
# print(f"{data['yoloid']} -> {data['status']}", end=' ')
else:
data['sequence'] = []
statuses = [(people[p]['yoloid'], people[p]['status']) for p in people]
# for id, status in statuses:
# print(f'{id}: {status}', end='\t')
# print()
frame_count += 1
cv2.imshow("YOLOv7 Pose Estimation Demo", im0)
key = cv2.waitKey(1) & 0xFF # Wait for 1 millisecond and get the pressed key
if key == ord('q'):
cv2.destroyAllWindows() # Close the window if 'q' is pressed
break
else:
break
cap.release()
class AutoPostureApp(App):
def __init__(self, opt):
super().__init__()
self.opt = opt
self.source = opt.source
self.device = opt.device
self.separation = opt.separation
self.length = opt.length
self.multiple = opt.multiple
def build(self):
# Create a layout for the GUI
layout = BoxLayout(orientation='vertical')
exit_button = Button(text='Exit Pose Estimation')
layout.add_widget(exit_button)
exit_button.bind(on_press=self.exit_app)
# Create a title label
title_label = Label(text='AutoPosture', font_size=20)
layout.add_widget(title_label)
# Create a checkbox to toggle the webcam display
self.webcam_checkbox = CheckBox()
self.webcam_checkbox.text = "View webcam"
layout.add_widget(self.webcam_checkbox)
# Create a button to start the pose estimation process
start_button = Button(text='Start')
layout.add_widget(start_button)
# self.start_button.bind(on_press=lambda instance: self.runPoseEstimation(self.opt))
start_button.bind(on_press=self.runPoseEstimation) # Bind the button to the action
return layout
def runPoseEstimation(self, instance):
run(self.source, self.opt.device, self.opt.separation, self.opt.length, self.opt.multiple)
def exit_app(self, instance):
App.get_running_app().stop()
def parse_opt():
parser = argparse.ArgumentParser()
parser.add_argument('-s', '--source', type=str, default='video/0', help='video/0 for webcam') #video source
parser.add_argument('-d', '--device', type=str, default='cpu', help='cpu/0,1,2,3(gpu)') #device arugments
parser.add_argument('-sep', '--separation', type=str, default='1', help='Each how many frames the prediction will be executed. Defaults to 1, increase for performance') #separation arugments
parser.add_argument('-l', '--length', type=str, default='10', help='Defines the length of the sequence. Defaults to 10, decrease for performance') #separation arugments
parser.add_argument('-m', '--multiple', default=False, action='store_true', help='Enable multiple-person detection') # Boolean for multiple person detection
opt = parser.parse_args()
return opt
def main(opt):
app = AutoPostureApp(opt)
app.run()
if __name__ == '__main__':
opt = parse_opt()
main(opt)
Datasets
Models
