PoseCorrection/backend/analyse_func.py

# -!- coding: utf-8 -!-
import numpy as np
import cv2 as cv
import math
import sys
from demo_func import run_openpose_for_image_front, run_openpose_for_image_side

def text_save(filename, data):
    file = open(filename, 'w', encoding='utf-8')
    for i in range(len(data)):
        s = str(data[i]).replace('[', '').replace(']', '')
        s = s.replace("'", '').replace(',', '') + '\n'
        file.write(s)
    file.close()
    print("保存成功")


def angle(v1, v2):
    dx1 = v1[2] - v1[0]
    dy1 = v1[3] - v1[1]
    dx2 = v2[2] - v2[0]
    dy2 = v2[3] - v2[1]
    angle1 = math.atan2(dy1, dx1)
    angle1 = round(angle1 * 180.0 / math.pi, 2)
    # print(angle1)
    angle2 = math.atan2(dy2, dx2)
    angle2 = round(angle2 * 180.0 / math.pi, 2)
    # print(angle2)
    if angle1 * angle2 >= 0:
        included_angle = abs(angle1 - angle2)
    else:
        included_angle = abs(angle1) + abs(angle2)
        if included_angle > 180:
            included_angle = 360 - included_angle
    return included_angle


def angle_for_xo(v1, v2):
    dx1 = v1[2] - v1[0]
    dy1 = v1[3] - v1[1]
    dx2 = v2[2] - v2[0]
    dy2 = v2[3] - v2[1]
    angle1 = math.atan2(dy1, dx1)
    angle1 = round(angle1 * 180.0 / math.pi, 2)
    angle2 = math.atan2(dy2, dx2)
    angle2 = round(angle2 * 180.0 / math.pi, 2)
    return angle1 - angle2


def cal_head_extra_stress(angle):
    HEAD_STRESS_ARGUMENT_0 = 11
    HEAD_STRESS_ARGUMENT_1 = 1.02777777777769
    HEAD_STRESS_ARGUMENT_2 = 0.00981481481482367
    HEAD_STRESS_ARGUMENT_3 = -0.000567901234568158
    HEAD_STRESS_ARGUMENT_4 = 0.00000576131687243021
    return (HEAD_STRESS_ARGUMENT_1 * angle + HEAD_STRESS_ARGUMENT_2 * pow(angle, 2) \
            + HEAD_STRESS_ARGUMENT_3 * pow(angle, 3) + HEAD_STRESS_ARGUMENT_4 * pow(angle, 4)) / HEAD_STRESS_ARGUMENT_0


def analyse_npy_front(npy_front, height):
    levels = {
        # (level, threshold)
        "head_angle": [
            (0, 4), (1, 10), (2, 20)
        ],
        "shoulder_angle": [
            (0, 2.8), (1, 8), (2, 15)
        ],
        "shoulder_offset": [
            (0, 3), (1, 8), (2, 15)
        ],
        "hip_angle": [
            (0, 2.19), (1, 8), (2, 15)
        ],
        "hip_offset": [
            (0, 2.19), (1, 8), (2, 15)
        ],
        "leg_angle": 4
    }

    result = []
    data = np.load(npy_front)

    print("------------------")
    result.append("------------------")

    print("正面分析内容如下：")
    result.append("正面分析内容如下：")

    # h-eye - h-elbow
    h_pixel_eye2elbow = abs(data[14][1] - data[3][1])
    # print("h_pixel{}".format(h_pixel_eye2elbow))

    h_real_eye2elbow = float(height) * 0.333
    # print("h_real{}".format(h_real_eye2elbow))

    # body mid x
    ankle_mid_x = (data[13][0] + data[10][0]) / 2

    # for head
    ear_r_x = data[16][0]
    ear_r_y = data[16][1]
    ear_l_x = data[17][0]
    ear_l_y = data[17][1]

    head_angle = angle([ear_r_x, ear_r_y, ear_l_x, ear_l_y], [0, 0, 1, 0])
    head_front_level = 0
    for level in levels["head_angle"]:
        head_front_level = level[0]
        if head_angle <= level[1]:
            break

    head_direction = None
    if ear_l_y < ear_l_y:
        head_direction = "左"
    elif ear_l_y > ear_r_y:
        head_direction = "右"
    if head_front_level > 0 and head_direction:
        print("头部分析：头部{}部高，向右倾斜，角度为{}度".format(head_direction, round(head_angle,2)))
        result.append("头部分析：头部{}部高，向右倾斜，角度为{}度".format(head_direction, round(head_angle,2)))
    else:
        print("头部分析：头部保持较好平衡")
        result.append("头部分析：头部保持较好平衡")
    result.append(head_front_level)
    print(head_front_level)

    # for shoulder
    shoulder_r_x = data[2][0]
    shoulder_r_y = data[2][1]
    shoulder_l_x = data[5][0]
    shoulder_l_y = data[5][1]
    shoulder_mid_x = (shoulder_l_x + shoulder_r_x) / 2
    neck_x = data[1][0] - 2
    shoulder_offset = (shoulder_mid_x - neck_x) * h_real_eye2elbow / h_pixel_eye2elbow
    shoulder_angle = angle([shoulder_r_x, shoulder_r_y, shoulder_l_x, shoulder_l_y], [0, 0, 1, 0])
    shoulder_angle = min(shoulder_angle, abs(shoulder_angle - 3))
    shoulder_angle_level = 0
    shoulder_offset_level = 0
    for level in levels["shoulder_angle"]:
        shoulder_angle_level = level[0]
        if shoulder_angle <= level[1]:
            break
    for level in levels["shoulder_offset"]:
        shoulder_offset_level = level[0]
        if shoulder_offset <= level[1]:
            break
    if shoulder_offset_level == 0:
        shoulder_level = shoulder_angle_level
    elif shoulder_angle_level == 0:
        shoulder_level = shoulder_offset_level
    else:
        shoulder_level = round(0.6 * shoulder_angle_level + 0.4 * shoulder_offset_level)

    result.append("肩部分析：")
    print("肩部分析：")
    shoulder_offset_direction = None
    if shoulder_offset > 0:
        shoulder_offset_direction = "左"
    elif shoulder_offset < 0:
        shoulder_offset_direction = "右"
    shoulder_angle_direction = None
    if shoulder_l_y < shoulder_r_y:
        shoulder_angle_direction = "右"
    elif shoulder_l_y > shoulder_r_y:
        shoulder_angle_direction = "左"
    if shoulder_offset_direction:
        print("肩部向{}横移，距离为{}cm".format(shoulder_offset_direction, round(abs(shoulder_offset), 2)))
        result.append("肩部向{}横移，距离为{}cm".format(shoulder_offset_direction, round(abs(shoulder_offset), 2)))
    else:
        print("肩部中心位居中轴，位置良好")
        result.append("肩部中心位居中轴，位置良好")

    if shoulder_angle_direction:
        print("肩部向{}倾斜角度为{}度".format(shoulder_angle_direction, round(shoulder_angle,2)))
        result.append("肩部向{}倾斜角度为{}度".format(shoulder_angle_direction, round(shoulder_angle,2)))
    else:
        print("肩部无左右倾斜")
        result.append("肩部无左右倾斜")
    result.append(shoulder_level)
    print(shoulder_level)

    # for hip
    hip_r_x = data[8][0]
    hip_r_y = data[8][1]
    hip_l_x = data[11][0]
    hip_l_y = data[11][1]
    hip_mid_x = (hip_r_x + hip_l_x) / 2
    hip_offset = (hip_mid_x - ankle_mid_x) * h_real_eye2elbow / h_pixel_eye2elbow
    hd_hip = abs(data[11][1] - data[8][1])
    d_hip = math.sqrt((hip_l_x - hip_r_x) ** 2 + (hip_l_y - hip_r_y) ** 2)
    rate_hip = hd_hip / d_hip
    hip_angle = angle([hip_r_x, hip_r_y, hip_l_x, hip_l_y], [0, 0, 1, 0])
    hip_angle_level = 0
    hip_offset_level = 0
    for level in levels["hip_angle"]:
        hip_angle_level = level[0]
        if hip_angle <= level[1]:
            break
    for level in levels["hip_offset"]:
        hip_offset_level = level[0]
        if hip_offset <= level[1]:
            break
    print("髋部分析：")
    result.append("髋部分析：")
    hip_offset_direction = None
    if hip_offset > 0:
        hip_offset_direction = "左"
    elif hip_offset < 0:
        hip_offset_direction = "右"
    hip_angle_direction = None
    if hip_l_y < hip_r_y:
        hip_angle_direction = "右"
    elif hip_l_y > hip_r_y:
        hip_angle_direction = "左"
    if hip_offset_direction:
        print("髋部向{}横移，距离为{}cm".format(hip_offset_direction, round(abs(hip_offset), 2)))
        result.append("髋部向{}横移，距离为{}cm".format(hip_offset_direction, round(abs(hip_offset), 2)))
    else:
        print("髋部中心位居中轴，位置良好")
        result.append("髋部中心位居中轴，位置良好")

    if hip_angle_direction:
        print("髋部向{}倾斜角度为{}度".format(hip_angle_direction, round(hip_angle,2)))
        result.append("髋部向{}倾斜角度为{}度".format(hip_angle_direction, round(hip_angle,2)))
    else:
        print("髋部角度保持较好平衡")
        result.append("髋部角度保持较好平衡")
    hip_level = round(0.6 * hip_offset_level + 0.4 * hip_angle_level)
    print(hip_level)
    result.append(hip_level)

    #

    print("腿型分析：")
    result.append("腿型分析：")

    # for knee distance
    knee_r_x = data[9][0]
    knee_r_y = data[9][1]
    knee_l_x = data[12][0]
    knee_l_y = data[12][1]

    ankle_r_x = data[10][0]
    ankle_r_y = data[10][1]
    ankle_l_x = data[13][0]
    ankle_l_y = data[13][1]

    xo_l_angle = angle_for_xo([hip_l_x, hip_l_y, knee_l_x, knee_l_y], [hip_l_x, hip_l_y, ankle_l_x, ankle_l_y])
    xo_r_angle = angle_for_xo([hip_r_x, hip_r_y, knee_r_x, knee_r_y], [hip_r_x, hip_r_y, ankle_r_x, ankle_r_y])
    xo_angle = max(abs(xo_l_angle), abs(xo_r_angle))
    if xo_angle > levels["leg_angle"]:
        xo_direction = "X" if xo_l_angle < 0 and xo_l_angle + xo_r_angle < 0 else "O"
        print("腿型可能为{}型腿".format(xo_direction))
        result.append("腿型可能为{}型腿".format(xo_direction))
    else:
        print("腿型基本正常，非X、O型腿")
        result.append("腿型基本正常，非X、O型腿")
    return result


def analyse_npy_side(npy_side, height):
    levels = {
        "head_forward_angle": [
            (0, 4.57), (1, 8), (2, 15)
        ],
        "up_forward_angle": [
            (0, 5), (1, 10), (2, 15)
        ],
        "hip_forward_offset": [
            (0, 3), (1, 5), (2, 8)
        ],
        "knee_forward_angle": [
            (0, 3.86), (1, 8), (2, 13)
        ]
    }

    result = []
    data = np.load(npy_side)
    h_pixel_eye2elbow = abs(data[14][1] - data[3][1])
    # print("h_pixel{}".format(h_pixel_eye2elbow))

    h_real_eye2elbow = float(height) * 0.333
    result.append(" ")
    print("------------------")
    result.append("------------------")
    print("侧面分析内容如下：")
    result.append("侧面分析内容如下：")

    # for head
    ear_r_x = data[len(data) - 1][0]
    ear_r_y = data[len(data) - 1][1]
    shoulder_r_x = data[2][0]
    shoulder_r_y = data[2][1]
    head_forward_angle = angle([shoulder_r_x, shoulder_r_y, ear_r_x, ear_r_y, ], [0, 0, 0, -1])
    head_forward_level = 0
    head_extra_stress = round(cal_head_extra_stress(head_forward_angle), 2) if 0 < head_forward_angle < 90 else 0
    for level in levels["head_forward_angle"]:
        head_forward_level = level[0]
        if head_forward_angle <= level[1]:
            break
    head_forward_direction = None
    if ear_r_x > shoulder_r_x:
        head_forward_direction = "前倾"
    elif ear_r_x < shoulder_r_x:
        head_forward_direction = "后仰"
    if head_forward_direction:
        print("头部{}角度为{}度，颈椎额外承受压力为头部重量{}倍".format(head_forward_direction, round(head_forward_angle, 2),
                                                   round(head_extra_stress,2)))
        result.append("头部{}角度为{}度，颈椎额外承受压力为头部重量{}倍".format(head_forward_direction, round(head_forward_angle, 2),
                                                           round(head_extra_stress,2)))
    else:
        print("无头部前倾问题")
        result.append("无头部前倾问题")
    print(head_forward_level)
    result.append(head_forward_level)

    # upper part of body
    neck_x = data[1][0]
    neck_y = data[1][1]

    ankle_r_x = data[10][0]
    ankle_r_y = data[10][1]
    up_angle = round(angle([ankle_r_x, ankle_r_y, neck_x, neck_y], [0, 0, 0, -1]), 2)
    up_angle = min(up_angle, abs(up_angle - 3))
    up_level = 0
    for level in levels["up_forward_angle"]:
        up_level = level[0]
        if up_angle <= level[1]:
            break
    up_forward_direction = "后仰" if neck_x < ankle_r_x else "前倾"
    up_state = "身体{}{}度".format(up_forward_direction, round(up_angle,2))
    print("站姿分析：{}".format(up_state))
    result.append("站姿分析：{}".format(up_state))
    print(up_level)
    result.append(up_level)

    # for hip
    hip_risk_level = 0
    hip_r_x = data[8][0] - 2
    hip_r_y = data[8][1]
    knee_r_x = data[9][0]
    knee_r_y = data[9][1]
    hip_offset = abs((hip_r_x - ankle_r_x)) * h_real_eye2elbow / h_pixel_eye2elbow
    for level in levels["hip_forward_offset"]:
        hip_risk_level = level[0]
        if hip_offset <= level[1]:
            break
    hip_forward_direction = None
    if hip_r_x > ankle_r_x:
        hip_forward_direction = "前移"
    elif hip_r_x < ankle_r_x:
        hip_forward_direction = "后移"
    if hip_forward_direction:
        hip_state = "髋关节{}{}cm".format(hip_forward_direction, round(hip_offset,2))
    else:
        hip_state = "髋关节正常"

    print("髋部分析：{}".format(hip_state))
    result.append("髋部分析：{}".format(hip_state))

    result.append(hip_risk_level)
    print(hip_risk_level)

    # for knee
    knee_angle_r = 180 - angle([hip_r_x, hip_r_y, knee_r_x, knee_r_y], [knee_r_x, knee_r_y, ankle_r_x, ankle_r_y])
    if knee_angle_r > 90:
        knee_angle_r = 180 - knee_angle_r
    knee_angle_r -= 3
    knee_forward_risk_level = 0
    for level in levels["knee_forward_angle"]:
        knee_forward_risk_level = level[0]
        if knee_angle_r <= level[1]:
            break

    if knee_forward_risk_level == 0:
        print("膝盖角度分析：膝盖无明显超伸")
        result.append("膝盖角度分析：膝盖无明显超伸")
    else:
        print("膝盖角度分析：下肢膝盖超伸角度为{}度".format(round(knee_angle_r, 2)))
        result.append("膝盖角度分析：下肢膝盖超伸角度为{}度".format(round(knee_angle_r, 2)))
    print(knee_forward_risk_level)
    result.append(knee_forward_risk_level)
    return result


def do_analysis(index, height):
    source_npy_front = "pose_temp_data/{}-0.npy".format(index)
    source_npy_side = "pose_temp_data/{}-1.npy".format(index)
    save_path = "pose_processed_txt/{}.txt".format(index)
    run_openpose_for_image_front(index)
    run_openpose_for_image_side(index)
    result_front = analyse_npy_front(source_npy_front, height)
    result_side = analyse_npy_side(source_npy_side, height)
    result = result_front + result_side
    text_save(save_path, result)


if __name__ == '__main__':
    heights = [104, 100, 104.5, 124, 107, 120.6, 111, 120, 118, 103.5, 101, 114, 128, 142, 117
        , 181, 105, 183, 150, 177
    ]
    for (i, h) in enumerate(heights):
        try:
            print("{}/{}".format(i + 1, len(heights)))
            do_analysis(i + 1, h)
        except Exception:
            print("\nWarning")
            print("照片{}有关键点未识别出来".format(i + 1))
            continue