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 src1.torch_openpose import torch_openpose
import src1
# 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')
    tp = torch_openpose('body_25')
    poses = tp(oriImg)
    np.save('pose_temp_data/{}-0.npy'.format(index), poses)
    canvas = copy.deepcopy(oriImg)
    canvas = src1.util.draw_bodypose(canvas, poses, 'body_25')
    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('pose_processed_images/{}-0-result.png'.format(index), canvas, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])


# 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)
#     print(candidate)
#     print("-----")
#     print(subset)
#
#     # --- 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
    # 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')
    tp = torch_openpose('body_25')
    print(tp)
    poses = tp(oriImg)
    np.save('pose_temp_data/{}-1.npy'.format(index), poses)
    canvas = copy.deepcopy(oriImg)
    canvas = src1.util.draw_bodypose(canvas, poses, 'body_25')
    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('pose_processed_images/{}-1-result.png'.format(index), canvas, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])


# 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):
    PNG_OR_NOT = True
    tp = torch_openpose('body_25')
    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)
    poses = tp(oriImg)
    np.save(p + '/capture_image{}-1.png'.format(index), poses)
    np.save(p + '/capture_image{}-1.png.npy'.format(index), poses)
    print(p + '/capture_image{}.png'.format(index))
    canvas = copy.deepcopy(oriImg)
    canvas = src1.util.draw_bodypose(canvas, poses, 'body_25')
    cv2.imwrite(p + '/capture_image_result{}.png'.format(index), canvas, [int(cv2.IMWRITE_PNG_COMPRESSION), 9])
    # data = np.load(p + '/capture_image{}-1.png.npy'.format(index))

    # data = np.load(p + '/capture_image{}-1.png.npy'.format(index))
    # print("25----")
    # print(data)
    #
    # plt.imshow(canvas[:, :, [2, 1, 0]])
    # plt.axis('off')
    # plt.savefig(p + '/capture_image_result{}.png'.format(index))
    # plt.show()

# def run_openpose_for_normal1(index,name):
#     print(index)
#
#     # 输入index和原来一样，就是序号，你可以更改一下输入或者输出，方便你们处理
#
#     # 判断是否是PNG格式，如果是，增添对alpha通道的处理
#     PNG_OR_NOT = True
#     # 调用的模型路径，无需更改
#     print('body_model')
#     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])
#     data = np.load(p + '/capture_image{}-1.png.npy'.format(index))
#     print("18----")
#     print(data)
#
#     plt.imshow(canvas[:, :, [2, 1, 0]])
#     plt.axis('off')
#     plt.savefig(p + '/capture_image_result{}.png'.format(index))
#     plt.show()


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