Permalink
Cannot retrieve contributors at this time
Name already in use
A tag already exists with the provided branch name. Many Git commands accept both tag and branch names, so creating this branch may cause unexpected behavior. Are you sure you want to create this branch?
Pedestrian_Direction_Analysis/img.py
Go to fileThis commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
175 lines (144 sloc)
8.72 KB
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import numpy as np | |
import cv2 | |
import os | |
import mediapipe as mp | |
import torch | |
from ultralytics import YOLO | |
def measure_horizontal_movement(left_heel, left_foot_index, right_heel, right_foot_index, arrow_scale=2): | |
# Calculate arrow parameters for left leg | |
left_arrow_length = int(np.sqrt((left_foot_index[0] - left_heel[0])**2 + (left_foot_index[1] - left_heel[1])**2) * arrow_scale) | |
if left_heel[0] < left_foot_index[0]: | |
left_arrow_tip = (int(left_heel[0] + left_arrow_length), int(left_heel[1])) | |
else: | |
left_arrow_tip = (int(left_heel[0] - left_arrow_length), int(left_heel[1])) | |
# Calculate arrow parameters for right leg | |
right_arrow_length = int(np.sqrt((right_foot_index[0] - right_heel[0])**2 + (right_foot_index[1] - right_heel[1])**2) * arrow_scale) | |
if right_heel[0] < right_foot_index[0]: | |
right_arrow_tip = (int(right_heel[0] + right_arrow_length), int(right_heel[1])) | |
else: | |
right_arrow_tip = (int(right_heel[0] - right_arrow_length), int(right_heel[1])) | |
# Determine movement direction based on arrow tips | |
if abs(left_arrow_tip[0] - left_heel[0])>20 or abs(right_arrow_tip[0] - right_heel[0])>20: | |
if left_arrow_tip[0] > left_heel[0] and right_arrow_tip[0] > right_heel[0]: # Both tips are to the right | |
return "Moving Right" | |
elif left_arrow_tip[0] < left_heel[0] and right_arrow_tip[0] < right_heel[0]: # Both tips are to the left | |
return"Moving Left" | |
else: | |
return "UP/DOWN" | |
def draw_arrow(image, heel, index, color=(0, 255, 0), thickness=2, arrow_scale=2, color_joint=(0, 0, 255)): | |
# Calculate arrow parameters | |
arrow_length = int(np.sqrt((index[0] - heel[0])**2 + (index[1] - heel[1])**2) * arrow_scale) | |
# Decide arrow direction based on the x-coordinates of heel and index points | |
if heel[0] < index[0]: | |
arrow_tip = (int(heel[0] + arrow_length), int(heel[1])) | |
else: | |
arrow_tip = (int(heel[0] - arrow_length), int(heel[1])) | |
# Draw the arrow | |
cv2.arrowedLine(image, (int(heel[0]), int(heel[1])), arrow_tip, color, thickness) | |
cv2.circle(image, (int(heel[0]), int(heel[1])), 3, color_joint, -1) # Draw a circle at the heel | |
cv2.circle(image, (int(index[0]), int(index[1])), 3, color_joint, -1) # Draw a circle at the index | |
# Function to add text on top of the person box | |
def add_label(image, bbox, text_lines, font_scale=0.5, font_thickness=1, text_color=(255, 255, 255), bg_color=(0, 0, 0)): | |
x1, y1, x2, y2 = bbox | |
# Calculate the size of the text box | |
(text_width, text_height), baseline = cv2.getTextSize(max(text_lines, key=len), cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness) | |
# Draw the background rectangle for the text box | |
text_height = text_height*len(text_lines) | |
cv2.rectangle(image, (x1, y1 - text_height - 10), (x1 + text_width, y1), bg_color, -1) | |
# Add text lines | |
y_text = y1 - 5 | |
for line in text_lines: | |
(text_width, text_height), baseline = cv2.getTextSize(line, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness) | |
cv2.putText(image, line, (x1, y_text), cv2.FONT_HERSHEY_SIMPLEX, font_scale, text_color, font_thickness, lineType=cv2.LINE_AA) | |
y_text -= int(text_height * 1.2) | |
# Load the models | |
model = YOLO('yolov8m-seg.pt') # load an official model | |
BaseOptions = mp.tasks.BaseOptions | |
PoseLandmarker = mp.tasks.vision.PoseLandmarker | |
PoseLandmarkerOptions = mp.tasks.vision.PoseLandmarkerOptions | |
VisionRunningMode = mp.tasks.vision.RunningMode | |
options = PoseLandmarkerOptions( | |
base_options=BaseOptions(model_asset_path='pose_landmarker_full.task'), | |
min_pose_presence_confidence = 0.0, | |
min_pose_detection_confidence = 0.0, | |
running_mode=VisionRunningMode.IMAGE) | |
def main(img_name,output): | |
# Predict with the model | |
results = model(img_name) # predict on an image | |
for result in results: | |
boxes = result.boxes | |
masks = result.masks | |
img = result.orig_img | |
for i, (mask, box) in enumerate(zip(masks,boxes)): | |
# check if person in the image | |
if box.cls!=0: | |
continue | |
# segmentatation of the person | |
xyxy = box.xyxy | |
orig_shape = box.orig_shape | |
mask_array = torch.any(mask.data, dim=0).int().cpu().detach().numpy() | |
# mask_array = mask.data.cpu().detach().numpy().transpose(1,2,0) | |
mask_array = cv2.resize(mask_array.astype(np.uint8),(orig_shape[1],orig_shape[0])) | |
mask_array = np.concatenate([np.expand_dims(mask_array,axis=2)]*3,axis=2) | |
segmented = img*mask_array | |
segmented[:,:,1:2] = segmented[:,:,1:2]*0 | |
img = cv2.addWeighted(img, 1, segmented, 0.8, 0) | |
xyxy = [int(j) for j in xyxy[0]] | |
cv2.rectangle(img, xyxy[:2],xyxy[2:], (255,0,0),2,2) | |
# person segmented to detect foot landmarks | |
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=segmented) | |
with PoseLandmarker.create_from_options(options) as landmarker: | |
# Perform pose landmarking on the provided single image. | |
# The pose landmarker must be created with the image mode. | |
pose_landmarker_result = landmarker.detect(mp_image) | |
if len(pose_landmarker_result.pose_landmarks): | |
joints = pose_landmarker_result.pose_landmarks[0] | |
# Left leg points | |
left_knee = joints[25] | |
left_ankle = joints[27] | |
left_heel = joints[29] | |
left_foot_index = joints[31] | |
left_leg = [[left_knee.x,left_knee.y], | |
[left_ankle.x,left_ankle.y], | |
[left_heel.x,left_heel.y], | |
[left_foot_index.x,left_foot_index.y] | |
] | |
# Right leg points | |
right_knee = joints[26] | |
right_ankle = joints[28] | |
right_heel = joints[30] | |
right_foot_index = joints[32] | |
right_leg = [[right_knee.x,right_knee.y], | |
[right_ankle.x,right_ankle.y], | |
[right_heel.x,right_heel.y], | |
[right_foot_index.x,right_foot_index.y] | |
] | |
# left_direction, right_direction = determine_foot_movement(left_leg, right_leg) | |
# print("Left foot direction:", left_direction) | |
# print("Right foot direction:", right_direction) | |
movement_direction = measure_horizontal_movement((left_heel.x*orig_shape[1],left_heel.y*orig_shape[0]), | |
(left_foot_index.x*orig_shape[1],left_foot_index.y*orig_shape[0]), | |
(right_heel.x*orig_shape[1] ,right_heel.y*orig_shape[0]), | |
(right_foot_index.x*orig_shape[1], right_foot_index.y*orig_shape[0]), | |
arrow_scale=2) | |
# movement_direction = measure_horizontal_movement(left_leg, right_leg, orig_shape[1]) | |
print("Horizontal Movement Direction:", movement_direction) | |
# print("Direction:", direction) | |
draw_arrow(img, (left_heel.x*orig_shape[1],left_heel.y*orig_shape[0]), | |
(left_foot_index.x*orig_shape[1],left_foot_index.y*orig_shape[0]), | |
color=(0, 255, 0), thickness=2, arrow_scale=2) | |
draw_arrow(img, (right_heel.x*orig_shape[1] ,right_heel.y*orig_shape[0]), | |
(right_foot_index.x*orig_shape[1], right_foot_index.y*orig_shape[0]), | |
color=(0, 255, 0), thickness=2, arrow_scale=2, | |
) | |
add_label(img, xyxy, [f'{movement_direction}' | |
]) | |
cv2.imwrite(os.path.join(output, os.path.basename(img_name)), img) | |
# Setup input and output directories and process each image | |
input_dir = 'input_img' | |
output_dir = 'output_img' | |
if not os.path.exists(output_dir): | |
os.makedirs(output_dir) | |
img_paths = [os.path.join(input_dir, f) for f in os.listdir(input_dir) if f.lower().endswith(('.jpg', '.jpeg', 'png'))] | |
for img_path in img_paths: | |
main(img_path, output_dir) |