Skip to content

Commit 709855a

Browse files
authored
added pose related api documentation (#129)
* added poses example code to readme * using urls in example to make it runnable * added comment
1 parent 8e04c88 commit 709855a

1 file changed

Lines changed: 32 additions & 5 deletions

File tree

README.md

Lines changed: 32 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -39,29 +39,56 @@ when using this project as a library:
3939
## Python API
4040

4141
```python
42+
from pathlib import Path
4243
import tensorflow as tf
4344
from tf_bodypix.api import download_model, load_model, BodyPixModelPaths
4445

46+
# setup input and output paths
47+
output_path = Path('./data/example-output')
48+
output_path.mkdir(parents=True, exist_ok=True)
49+
input_url = (
50+
'https://www.dropbox.com/s/7tsaqgdp149d8aj/serious-black-businesswoman-sitting-at-desk-in-office-5669603.jpg?dl=1'
51+
)
52+
local_input_path = tf.keras.utils.get_file(origin=input_url)
53+
54+
# load model (once)
4555
bodypix_model = load_model(download_model(
4656
BodyPixModelPaths.MOBILENET_FLOAT_50_STRIDE_16
4757
))
4858

49-
image = tf.keras.preprocessing.image.load_img(
50-
'/path/to/input-image.jpg'
51-
)
59+
# get prediction result
60+
image = tf.keras.preprocessing.image.load_img(local_input_path)
5261
image_array = tf.keras.preprocessing.image.img_to_array(image)
5362
result = bodypix_model.predict_single(image_array)
63+
64+
# simple mask
5465
mask = result.get_mask(threshold=0.75)
5566
tf.keras.preprocessing.image.save_img(
56-
'/path/to/output-mask.jpg',
67+
f'{output_path}/output-mask.jpg',
5768
mask
5869
)
5970

71+
# colored mask (separate colour for each body part)
6072
colored_mask = result.get_colored_part_mask(mask)
6173
tf.keras.preprocessing.image.save_img(
62-
'/path/to/output-colored-mask.jpg',
74+
f'{output_path}/output-colored-mask.jpg',
6375
colored_mask
6476
)
77+
78+
# poses
79+
from tf_bodypix.draw import draw_poses # utility function using OpenCV
80+
81+
poses = result.get_poses()
82+
image_with_poses = draw_poses(
83+
image_array.copy(), # create a copy to ensure we are not modifing the source image
84+
poses,
85+
keypoints_color=(255, 100, 100),
86+
skeleton_color=(100, 100, 255)
87+
)
88+
tf.keras.preprocessing.image.save_img(
89+
f'{output_path}/output-poses.jpg',
90+
image_with_poses
91+
)
6592
```
6693

6794
## CLI

0 commit comments

Comments
 (0)