PoseCorrection/backend/demo_func.py

import cv2
import matplotlib.pyplot as plt
import copy
import numpy as np
from PIL import Image
import os

from src import util
from src.body import Body
# from src.infer_video_d2 import infer_image,load_predictor
from video_pose.inference.infer_video_d2 import run_3d
from video_pose.data.prepare_data_2d_custom import custom
from video_pose.run import run_video_reconstruction




def run_3d_for_video(original_video_path):
    run_3d(original_video_path)
    video_analyzed_path = original_video_path + ".npz"
    print(video_analyzed_path)
    custom(video_analyzed_path)
    compressed_analyzed_path = video_analyzed_path.replace(".npz", "_compressed.npz")
    run_video_reconstruction(original_video_path, compressed_analyzed_path)
    reconstructed_video_path = original_video_path.replace(".mp4", "_reconstructed.mp4")
    print(reconstructed_video_path)
    # return reconstructed_video_path


def run_openpose_for_image_front(index):
    PNG_OR_NOT = True
    # hand_estimation = Hand('model/hand_pose_model.pth')
    test_image = 'pose_source_images/{}-0.png'.format(index)
    if os.path.exists(test_image) is False:
        PNG_OR_NOT = False
        test_image = 'pose_source_images/{}-0.jpg'.format(index)
    oriImg = cv2.imread(test_image)  # B,G,R order
    im = cv2.imread(test_image, cv2.IMREAD_UNCHANGED)

    # --- for openpose ---
    body_estimation = Body('model/body_pose_model.pth')
    candidate, subset = body_estimation(oriImg)

    # --- for detectron ---
    # predictor = load_predictor()
    # candidate, subset = infer_image(oriImg, predictor)
    np.save('pose_temp_data/{}-0.npy'.format(index), candidate)

    canvas = copy.deepcopy(oriImg)
    canvas = util.draw_bodypose(canvas, candidate, subset)

    # sssssssssssss
    if PNG_OR_NOT == True:
        alpha_value_new = np.reshape(im[:, :, 3], (im.shape[0], im.shape[1], 1))
        canvas = np.c_[canvas, alpha_value_new]

    # test = np.c_[im[:,:,0:3],alpha_value_new]
    # print(test)
    # print(im)

    # cv2.imwrite('pose_processed_images/tmp_transparent_ori.png', oriImg, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    cv2.imwrite('pose_processed_images/{}-0-result.png'.format(index), canvas, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    # cv2.imwrite('pose_processed_images/tmp_transparent_test.png', test, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    # cv2.imwrite('pose_processed_images/tmp_transparent_im.png', im, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])

    # plt.imshow(canvas[:, :, [2, 1, 0]])
    # plt.axis('off')
    # plt.savefig('pose_processed_images/{}-1-result.png'.format(index))
    # plt.show()


def run_openpose_for_image_side(index):
    PNG_OR_NOT = True
    body_estimation = Body('model/body_pose_model.pth')
    # hand_estimation = Hand('model/hand_pose_model.pth')
    test_image = 'pose_source_images/{}-1.png'.format(index)
    if os.path.exists(test_image) is False:
        PNG_OR_NOT = False
        test_image = 'pose_source_images/{}-1.jpg'.format(index)
    oriImg = cv2.imread(test_image)  # B,G,R order
    im = cv2.imread(test_image, cv2.IMREAD_UNCHANGED)
    # --- for openpose ---
    body_estimation = Body('model/body_pose_model.pth')
    candidate, subset = body_estimation(oriImg)

    # --- for detectron ---
    # predictor = load_predictor()
    # candidate, subset = infer_image(oriImg, predictor)
    np.save('pose_temp_data/{}-1.npy'.format(index), candidate)

    canvas = copy.deepcopy(oriImg)
    canvas = util.draw_bodypose(canvas, candidate, subset)

    if PNG_OR_NOT == True:
        alpha_value_new = np.reshape(im[:, :, 3], (im.shape[0], im.shape[1], 1))
        canvas = np.c_[canvas, alpha_value_new]
    # test = np.c_[im[:,:,0:3],alpha_value_new]
    # print(test)
    # print(im)

    # cv2.imwrite('pose_processed_images/tmp_transparent_ori.png', oriImg, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    cv2.imwrite('pose_processed_images/{}-1-result.png'.format(index), canvas, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    # cv2.imwrite('pose_processed_images/tmp_transparent_test.png', test, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    # cv2.imwrite('pose_processed_images/tmp_transparent_im.png', im, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])

    # plt.imshow(canvas[:, :, [2, 1, 0]])
    # plt.axis('off')
    # plt.savefig('pose_processed_images/{}-1-result.png'.format(index))
    # plt.show()


def run_openpose_for_normal(index,name):
    print(index)
    
    # 输入index和原来一样，就是序号，你可以更改一下输入或者输出，方便你们处理

    # 判断是否是PNG格式，如果是，增添对alpha通道的处理
    PNG_OR_NOT = True
    # 调用的模型路径，无需更改
    body_estimation = Body('model/body_pose_model.pth')
    #原图像的位置，需要更改 ！！
    p='./capture_image/'+name
    test_image = p+'/capture_image{}.png'.format(index)
    if os.path.exists(test_image) is False:
        PNG_OR_NOT = False
        # 原图像的位置，需要更改 ！！
        test_image = p+'/capture_image{}.jpg'.format(index)
    oriImg = cv2.imread(test_image)  # B,G,R order
    im = cv2.imread(test_image, cv2.IMREAD_UNCHANGED)

    candidate, subset = body_estimation(oriImg)
    # 坐标点数值保存路径，需要更改！！
    np.save(p+'/capture_image{}-1.png'.format(index), candidate)
    print(p+'/capture_image{}.png'.format(index))

    # 处理
    canvas = copy.deepcopy(oriImg)
    canvas = util.draw_bodypose(canvas, candidate, subset)
    #if PNG_OR_NOT == True:
        #alpha_value_new = np.reshape(im[:,:,3],(im.shape[0],im.shape[1],1))
        #canvas = np.c_[canvas, alpha_value_new]

    # 结果图片的保存路径，需要更改 ！！
    cv2.imwrite(p+'/capture_image_result{}.png'.format(index), canvas, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])


if __name__ == '__main__':
    # run_openpose_for_normal(1)
    input_video = "video/15.mp4"
    run_3d_for_video(input_video)