Skip to content

Commit e00377a

Browse files
alex9311alex9311
authored
add demo code (leoxiaobin#161)
* add demo code * demo code review update: use BGR for pose model * demo code review update 2
1 parent 163eb47 commit e00377a

7 files changed

+583
-0
lines changed

‎demo/.gitignore

+3
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
output
2+
models
3+
videos

‎demo/Dockerfile

+112
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
FROM nvidia/cuda:10.2-cudnn7-devel-ubuntu16.04
2+
3+
ENV OPENCV_VERSION="3.4.6"
4+
5+
# Basic toolchain
6+
RUN apt-get update && apt-get install -y \
7+
apt-utils \
8+
build-essential \
9+
git \
10+
wget \
11+
unzip \
12+
yasm \
13+
pkg-config \
14+
libcurl4-openssl-dev \
15+
zlib1g-dev \
16+
htop \
17+
cmake \
18+
nano \
19+
python3-pip \
20+
python3-dev \
21+
python3-tk \
22+
libx264-dev \
23+
&& cd /usr/local/bin \
24+
&& ln -s /usr/bin/python3 python \
25+
&& pip3 install --upgrade pip \
26+
&& apt-get autoremove -y
27+
28+
# Getting OpenCV dependencies available with apt
29+
RUN apt-get update && apt-get install -y \
30+
libeigen3-dev \
31+
libjpeg-dev \
32+
libpng-dev \
33+
libtiff-dev \
34+
libjasper-dev \
35+
libswscale-dev \
36+
libavcodec-dev \
37+
libavformat-dev && \
38+
apt-get autoremove -y
39+
40+
# Getting other dependencies
41+
RUN apt-get update && apt-get install -y \
42+
cppcheck \
43+
graphviz \
44+
doxygen \
45+
p7zip-full \
46+
libdlib18 \
47+
libdlib-dev && \
48+
apt-get autoremove -y
49+
50+
51+
# Install OpenCV + OpenCV contrib (takes forever)
52+
RUN mkdir -p /tmp && \
53+
cd /tmp && \
54+
wget --no-check-certificate -O opencv.zip https://github.com/opencv/opencv/archive/${OPENCV_VERSION}.zip && \
55+
wget --no-check-certificate -O opencv_contrib.zip https://github.com/opencv/opencv_contrib/archive/${OPENCV_VERSION}.zip && \
56+
unzip opencv.zip && \
57+
unzip opencv_contrib.zip && \
58+
mkdir opencv-${OPENCV_VERSION}/build && \
59+
cd opencv-${OPENCV_VERSION}/build && \
60+
cmake -D CMAKE_BUILD_TYPE=RELEASE \
61+
-D CMAKE_INSTALL_PREFIX=/usr/local \
62+
-D WITH_CUDA=ON \
63+
-D CUDA_FAST_MATH=1 \
64+
-D WITH_CUBLAS=1 \
65+
-D WITH_FFMPEG=ON \
66+
-D WITH_OPENCL=ON \
67+
-D WITH_V4L=ON \
68+
-D WITH_OPENGL=ON \
69+
-D OPENCV_EXTRA_MODULES_PATH=/tmp/opencv_contrib-${OPENCV_VERSION}/modules \
70+
.. && \
71+
make -j$(nproc) && \
72+
make install && \
73+
echo "/usr/local/lib" > /etc/ld.so.conf.d/opencv.conf && \
74+
ldconfig && \
75+
cd /tmp && \
76+
rm -rf opencv-${OPENCV_VERSION} opencv.zip opencv_contrib-${OPENCV_VERSION} opencv_contrib.zip && \
77+
cd /
78+
79+
# Compile and install ffmpeg from source
80+
RUN git clone https://github.com/FFmpeg/FFmpeg /root/ffmpeg && \
81+
cd /root/ffmpeg && \
82+
./configure --enable-gpl --enable-libx264 --enable-nonfree --disable-shared --extra-cflags=-I/usr/local/include && \
83+
make -j8 && make install -j8
84+
85+
# clone deep-high-resolution-net
86+
ARG POSE_ROOT=/pose_root
87+
RUN git clone https://github.com/leoxiaobin/deep-high-resolution-net.pytorch.git $POSE_ROOT
88+
WORKDIR $POSE_ROOT
89+
RUN mkdir output && mkdir log
90+
91+
RUN pip3 install -r requirements.txt && \
92+
pip3 install torch==1.1.0 \
93+
torchvision==0.3.0 \
94+
opencv-python \
95+
pillow==6.2.1
96+
97+
# build deep-high-resolution-net lib
98+
WORKDIR $POSE_ROOT/lib
99+
RUN make
100+
101+
# install COCO API
102+
ARG COCOAPI=/cocoapi
103+
RUN git clone https://github.com/cocodataset/cocoapi.git $COCOAPI
104+
WORKDIR $COCOAPI/PythonAPI
105+
# Install into global site-packages
106+
RUN make install
107+
108+
# download fastrrnn pretrained model for person detection
109+
RUN python -c "import torchvision; model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True); model.eval()"
110+
111+
COPY inference.py $POSE_ROOT/tools
112+
COPY inference-config.yaml $POSE_ROOT/

‎demo/README.md

+41
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
This demo code is meant to be run on a video and includes a person detector.
2+
[Nvidia-docker](https://github.com/NVIDIA/nvidia-docker) and GPUs are required.
3+
It only expects there to be one person in each frame of video, though the code could easily be extended to support multiple people.
4+
5+
### Prep
6+
1. Download the researchers' pretrained pose estimator from [google drive](https://drive.google.com/drive/folders/1hOTihvbyIxsm5ygDpbUuJ7O_tzv4oXjC?usp=sharing) to this directory under `models/`
7+
2. Put the video file you'd like to infer on in this directory under `videos`
8+
3. build the docker container in this directory with `./build-docker.sh` (this can take time because it involves compiling opencv)
9+
4. update the `inference-config.yaml` file to reflect the number of GPUs you have available
10+
11+
### Running the Model
12+
Start your docker container with:
13+
```
14+
nvidia-docker run --rm -it \
15+
-v $(pwd)/output:/output \
16+
-v $(pwd)/videos:/videos \
17+
-v $(pwd)/models:/models \
18+
-w /pose_root \
19+
hrnet_demo_inference \
20+
/bin/bash
21+
```
22+
23+
Once the container is running, you can run inference with:
24+
```
25+
python tools/inference.py \
26+
--cfg inference-config.yaml \
27+
--videoFile /videos/my-video.mp4 \
28+
--inferenceFps 10 \
29+
--writeBoxFrames \
30+
TEST.MODEL_FILE \
31+
/models/pytorch/pose_coco/pose_hrnet_w32_384x288.pth
32+
```
33+
34+
The command above will output frames with boxes,
35+
frames with poses,
36+
a video with poses,
37+
and a csv with the keypoint coordinates for each frame.
38+
39+
![](hrnet-demo.gif)
40+
41+
Original source for demo video above is licensed for `Free for commercial use No attribution required` by [Pixabay](https://pixabay.com/service/license/)

‎demo/build-docker.sh

+1
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
docker build -t hrnet_demo_inference .

‎demo/hrnet-demo.gif

2.17 MB
Loading

‎demo/inference-config.yaml

+127
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
AUTO_RESUME: true
2+
CUDNN:
3+
BENCHMARK: true
4+
DETERMINISTIC: false
5+
ENABLED: true
6+
DATA_DIR: ''
7+
GPUS: (0,)
8+
OUTPUT_DIR: 'output'
9+
LOG_DIR: 'log'
10+
WORKERS: 24
11+
PRINT_FREQ: 100
12+
13+
DATASET:
14+
COLOR_RGB: true
15+
DATASET: 'coco'
16+
DATA_FORMAT: jpg
17+
FLIP: true
18+
NUM_JOINTS_HALF_BODY: 8
19+
PROB_HALF_BODY: 0.3
20+
ROOT: 'data/coco/'
21+
ROT_FACTOR: 45
22+
SCALE_FACTOR: 0.35
23+
TEST_SET: 'val2017'
24+
TRAIN_SET: 'train2017'
25+
MODEL:
26+
INIT_WEIGHTS: true
27+
NAME: pose_hrnet
28+
NUM_JOINTS: 17
29+
PRETRAINED: 'models/pytorch/imagenet/hrnet_w32-36af842e.pth'
30+
TARGET_TYPE: gaussian
31+
IMAGE_SIZE:
32+
- 288
33+
- 384
34+
HEATMAP_SIZE:
35+
- 72
36+
- 96
37+
SIGMA: 3
38+
EXTRA:
39+
PRETRAINED_LAYERS:
40+
- 'conv1'
41+
- 'bn1'
42+
- 'conv2'
43+
- 'bn2'
44+
- 'layer1'
45+
- 'transition1'
46+
- 'stage2'
47+
- 'transition2'
48+
- 'stage3'
49+
- 'transition3'
50+
- 'stage4'
51+
FINAL_CONV_KERNEL: 1
52+
STAGE2:
53+
NUM_MODULES: 1
54+
NUM_BRANCHES: 2
55+
BLOCK: BASIC
56+
NUM_BLOCKS:
57+
- 4
58+
- 4
59+
NUM_CHANNELS:
60+
- 32
61+
- 64
62+
FUSE_METHOD: SUM
63+
STAGE3:
64+
NUM_MODULES: 4
65+
NUM_BRANCHES: 3
66+
BLOCK: BASIC
67+
NUM_BLOCKS:
68+
- 4
69+
- 4
70+
- 4
71+
NUM_CHANNELS:
72+
- 32
73+
- 64
74+
- 128
75+
FUSE_METHOD: SUM
76+
STAGE4:
77+
NUM_MODULES: 3
78+
NUM_BRANCHES: 4
79+
BLOCK: BASIC
80+
NUM_BLOCKS:
81+
- 4
82+
- 4
83+
- 4
84+
- 4
85+
NUM_CHANNELS:
86+
- 32
87+
- 64
88+
- 128
89+
- 256
90+
FUSE_METHOD: SUM
91+
LOSS:
92+
USE_TARGET_WEIGHT: true
93+
TRAIN:
94+
BATCH_SIZE_PER_GPU: 32
95+
SHUFFLE: true
96+
BEGIN_EPOCH: 0
97+
END_EPOCH: 210
98+
OPTIMIZER: adam
99+
LR: 0.001
100+
LR_FACTOR: 0.1
101+
LR_STEP:
102+
- 170
103+
- 200
104+
WD: 0.0001
105+
GAMMA1: 0.99
106+
GAMMA2: 0.0
107+
MOMENTUM: 0.9
108+
NESTEROV: false
109+
TEST:
110+
BATCH_SIZE_PER_GPU: 32
111+
COCO_BBOX_FILE: 'data/coco/person_detection_results/COCO_val2017_detections_AP_H_56_person.json'
112+
BBOX_THRE: 1.0
113+
IMAGE_THRE: 0.0
114+
IN_VIS_THRE: 0.2
115+
MODEL_FILE: ''
116+
NMS_THRE: 1.0
117+
OKS_THRE: 0.9
118+
USE_GT_BBOX: true
119+
FLIP_TEST: true
120+
POST_PROCESS: true
121+
SHIFT_HEATMAP: true
122+
DEBUG:
123+
DEBUG: true
124+
SAVE_BATCH_IMAGES_GT: true
125+
SAVE_BATCH_IMAGES_PRED: true
126+
SAVE_HEATMAPS_GT: true
127+
SAVE_HEATMAPS_PRED: true

0 commit comments

Comments
 (0)