img.py

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)