Merge pull request Hzzone#23 from FabianHertwig/master

Add section to readme how to setup the environment and improve imports

README.md

@@ -16,11 +16,41 @@ If anybody wants a pure python wrapper, please refer to my [pytorch implementati
 
 Don't be mean to star this repo if it helps your research.
 
-### Model Download
+### Getting Started
+
+#### Install Requriements
+
+Create a python 3.7 environement, eg:
+
+    conda create -n pytorch-openpose python=3.7
+    conda activate pytorch-openpose
+
+Install pytorch by following the quick start guide here (use pip) https://download.pytorch.org/whl/torch_stable.html
+
+Install other requirements with pip
+
+    pip install -r requirements.txt
+
+#### Download the Models
+
 * [dropbox](https://www.dropbox.com/sh/7xbup2qsn7vvjxo/AABWFksdlgOMXR_r5v3RwKRYa?dl=0)
 
 `*.pth` files are pytorch model, you could also download caffemodel file if you want to use caffe as backend.
 
+Download the pytorch models and put them in a directory named `model` in the project root directory
+
+#### Run the Demo
+
+Run:
+
+    python demo_camera.py
+
+to run a demo with a feed from your webcam or run
+
+    python demo.py
+
+to use a image from the images folder.
+
 ### Todo list
 - [x] convert caffemodel to pytorch.
 - [x] Body Pose Estimation.

demo.py

@@ -1,14 +1,13 @@
-import sys
-sys.path.insert(0, 'python')
 import cv2
-import model
-import util
-from hand import Hand
-from body import Body
 import matplotlib.pyplot as plt
 import copy
 import numpy as np
 
+from src import model
+from src import util
+from src.body import Body
+from src.hand import Hand
+
 body_estimation = Body('model/body_pose_model.pth')
 hand_estimation = Hand('model/hand_pose_model.pth')
 

demo_camera.py

@@ -1,17 +1,19 @@
-import sys
-sys.path.insert(0, 'python')
 import cv2
-import model
-import util
-from hand import Hand
-from body import Body
 import matplotlib.pyplot as plt
 import copy
 import numpy as np
+import torch
+
+from src import model
+from src import util
+from src.body import Body
+from src.hand import Hand
 
 body_estimation = Body('model/body_pose_model.pth')
 hand_estimation = Hand('model/hand_pose_model.pth')
 
+print(f"Torch device: {torch.cuda.get_device_name()}")
+
 cap = cv2.VideoCapture(0)
 cap.set(3, 640)
 cap.set(4, 480)

requirements.txt

@@ -0,0 +1,6 @@
+numpy
+matplotlib
+opencv-python
+scipy
+scikit-image
+tqdm

python/body.py → src/body.py

@@ -1,15 +1,16 @@
 import cv2
 import numpy as np
-import util
 import math
 import time
 from scipy.ndimage.filters import gaussian_filter
 import matplotlib.pyplot as plt
 import matplotlib
-from model import bodypose_model
 import torch
 from torchvision import transforms
 
+from src import util
+from src.model import bodypose_model
+
 class Body(object):
     def __init__(self, model_path):
         self.model = bodypose_model()

python/hand.py → src/hand.py

@@ -1,16 +1,17 @@
 import cv2
 import json
 import numpy as np
-import util
 import math
 import time
 from scipy.ndimage.filters import gaussian_filter
 import matplotlib.pyplot as plt
 import matplotlib
-from model import handpose_model
 import torch
 from skimage.measure import label
 
+from src.model import handpose_model
+from src import util
+
 class Hand(object):
     def __init__(self, model_path):
         self.model = handpose_model()

python/hand_model_outputsize.py → src/hand_model_outputsize.py

@@ -1,8 +1,9 @@
-from model import handpose_model
 import torch
 from tqdm import tqdm
 import json
 
+from src.model import handpose_model
+
 model = handpose_model()
 
 size = {}