@@ -39,29 +39,56 @@ when using this project as a library:
3939## Python API
4040
4141``` python
42+ from pathlib import Path
4243import tensorflow as tf
4344from 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)
4555bodypix_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)
5261image_array = tf.keras.preprocessing.image.img_to_array(image)
5362result = bodypix_model.predict_single(image_array)
63+
64+ # simple mask
5465mask = result.get_mask(threshold = 0.75 )
5566tf.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)
6072colored_mask = result.get_colored_part_mask(mask)
6173tf.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