import argparse
import os
import sys
import time
import cv2
import matplotlib.pyplot as plt
import numpy as np
import requests
from internal.prediction_client import predict_http_request
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
POSEWEIGHTS = 'src_models/yolov7-w6-pose.pt'
@torch.no_grad()
def run(source, device, separation, length, multiple):
# 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 = []
current_score = 0
current_status = 'good'
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()
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):
run(**vars(opt))
if __name__ == "__main__":
opt = parse_opt()
strip_optimizer(opt.device, POSEWEIGHTS)
main(opt)
Datasets
Models
