Modified to use dnn module of opencv instead caffe

predict.py

@@ -2,16 +2,19 @@
 # OpenPose Position to 3D Position
 # ----------------------------------------------
 
-import cv2
+import cv2 as cv
 import sys
 import numpy as np
 import pandas as pd
 import os
-import caffe
+import argparse
 import h5py
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 
+import plaidml.keras
+plaidml.keras.install_backend()
+
 from mpl_toolkits.mplot3d import Axes3D
 from keras.models import load_model
 
@@ -20,9 +23,9 @@
 # ----------------------------------------------
 
 MODEL_HDF5 = "3d-pose-baseline.hdf5"
-IMAGE_PATH = "images/running.jpg"
-CAFFE_MODEL = 'pose_iter_440000.caffemodel'
-PROTOTXT = 'pose_deploy.prototxt'
+#IMAGE_PATH = "images/running.jpg"
+#CAFFE_MODEL = 'pose_iter_440000.caffemodel'
+#PROTOTXT = 'pose_deploy.prototxt'
 
 # ----------------------------------------------
 # Input Data
@@ -37,14 +40,21 @@
   IMAGE_WIDTH = int(sys.argv[3])
   IMAGE_HEIGHT = int(sys.argv[4])
 
-if not os.path.exists(IMAGE_PATH):
-	print(IMAGE_PATH+" not found")
+parser = argparse.ArgumentParser()
+parser.add_argument('--input', type=str, default='images/running.jpg')
+parser.add_argument('--width', type=int, default=368)
+parser.add_argument('--height', type=int, default=368)
+parser.add_argument('--proto', type=str, default='pose_deploy.prototxt')
+parser.add_argument('--model', type=str, default='pose_iter_440000.caffemodel')
+args = parser.parse_args()
+
+if not os.path.exists(args.input):
+	print(args.input + " not found")
 	sys.exit(1)
 
-print IMAGE_PATH
-input_img = cv2.imread(IMAGE_PATH)
+input_img = cv.imread(args.input)
 
-img = cv2.resize(input_img, (IMAGE_WIDTH, IMAGE_HEIGHT))
+img = cv.resize(input_img, (args.width, args.height))
 
 img = img[...,::-1]  #BGR 2 RGB
 
@@ -53,13 +63,12 @@
 
 data = data / 255.0
 
-net  = caffe.Net(PROTOTXT, CAFFE_MODEL, caffe.TEST)
+net  = cv.dnn.readNetFromCaffe(args.proto, args.model) #caffe.Net(PROTOTXT, CAFFE_MODEL, caffe.TEST)
 data = data.transpose((0, 3, 1, 2))
-out = net.forward_all(data = data)
-
-paf = out[net.outputs[0]]
-confidence = out[net.outputs[1]]
+net.setInput(data)
+out = net.forward()
 
+confidence = out
 # ----------------------------------------------
 # Display output
 # ----------------------------------------------
@@ -69,11 +78,11 @@
 
 for i in range(confidence.shape[1]):
 	probMap = confidence[0, i, :, :]
-	minVal, prob, minLoc, point = cv2.minMaxLoc(probMap)
+	minVal, prob, minLoc, point = cv.minMaxLoc(probMap)
 
 	x = (input_img.shape[1] * point[0]) / confidence.shape[3]
 	y = (input_img.shape[0] * point[1]) / confidence.shape[2] 
- 
+
 	if prob > threshold : 
 		points.append(x)
 		points.append(y)

train.py

@@ -6,7 +6,10 @@
 import numpy as np
 import h5py
 
-os.environ['KERAS_BACKEND'] = 'tensorflow'
+import plaidml.keras
+plaidml.keras.install_backend()
+
+#os.environ['KERAS_BACKEND'] = 'tensorflow'
 
 from keras.layers.convolutional import Convolution2D, Conv2D
 from keras.layers.convolutional import MaxPooling2D