08_heatmaps_from_image.py 3.2 KB
Newer Older
R
Raaj 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# From Python
# It requires OpenCV installed for Python
import sys
import cv2
import os
from sys import platform
import argparse

# Import Openpose (Windows/Ubuntu/OSX)
dir_path = os.path.dirname(os.path.realpath(__file__))
try:
    # Windows Import
    if platform == "win32":
        # Change these variables to point to the correct folder (Release/x64 etc.) 
        sys.path.append(dir_path + '/../../python/openpose/Release');
        os.environ['PATH']  = os.environ['PATH'] + ';' + dir_path + '/../../x64/Release;' +  dir_path + '/../../bin;'
17
        import pyopenpose as op
R
Raaj 已提交
18 19 20 21 22
    else:
        # Change these variables to point to the correct folder (Release/x64 etc.) 
        sys.path.append('../../python');
        # If you run `make install` (default path is `/usr/local/python` for Ubuntu), you can also access the OpenPose/python module from there. This will install OpenPose and the python library at your desired installation path. Ensure that this is in your python path in order to use it.
        # sys.path.append('/usr/local/python')
23
        from openpose import pyopenpose as op
24 25 26
except ImportError as e:
    print('Error: OpenPose library could not be found. Did you enable `BUILD_PYTHON` in CMake and have this Python script in the right folder?')
    raise e
R
Raaj 已提交
27 28 29 30 31 32 33 34 35 36

# Flags
parser = argparse.ArgumentParser()
parser.add_argument("--image_path", default="../../../examples/media/COCO_val2014_000000000192.jpg", help="Process an image. Read all standard formats (jpg, png, bmp, etc.).")
args = parser.parse_known_args()

# Custom Params (refer to include/openpose/flags.hpp for more parameters)
params = dict()
params["model_folder"] = "../../../models/"
params["heatmaps_add_parts"] = True
R
Raaj 已提交
37
params["heatmaps_add_bkg"] = True
R
Raaj 已提交
38
params["heatmaps_add_PAFs"] = True
R
Raaj 已提交
39
params["heatmaps_scale"] = 2
R
Raaj 已提交
40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72

# Add others in path?
for i in range(0, len(args[1])):
    curr_item = args[1][i]
    if i != len(args[1])-1: next_item = args[1][i+1]
    else: next_item = "1"
    if "--" in curr_item and "--" in next_item:
        key = curr_item.replace('-','')
        if key not in params:  params[key] = "1"
    elif "--" in curr_item and "--" not in next_item:
        key = curr_item.replace('-','')
        if key not in params: params[key] = next_item

# Construct it from system arguments
# op.init_argv(args[1])
# oppython = op.OpenposePython()

# Starting OpenPose
opWrapper = op.WrapperPython()
opWrapper.configure(params)
opWrapper.start()

# Process Image
datum = op.Datum()
imageToProcess = cv2.imread(args[0].image_path)
datum.cvInputData = imageToProcess
opWrapper.emplaceAndPop([datum])

# Process outputs
outputImageF = (datum.inputNetData[0].copy())[0,:,:,:] + 0.5
outputImageF = cv2.merge([outputImageF[0,:,:], outputImageF[1,:,:], outputImageF[2,:,:]])
outputImageF = (outputImageF*255.).astype(dtype='uint8')
heatmaps = datum.poseHeatMaps.copy()
R
Raaj 已提交
73
heatmaps = (heatmaps).astype(dtype='uint8')
R
Raaj 已提交
74 75 76 77 78 79 80 81

# Display Image
counter = 0
while 1:
    num_maps = heatmaps.shape[0]
    heatmap = heatmaps[counter, :, :].copy()
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    combined = cv2.addWeighted(outputImageF, 0.5, heatmap, 0.5, 0)
R
Raaj 已提交
82 83 84 85
    cv2.imshow("OpenPose 1.4.0 - Tutorial Python API", combined)
    key = cv2.waitKey(-1)
    if key == 27:
        break
R
Raaj 已提交
86 87
    counter += 1
    counter = counter % num_maps