Merge pull request isarandi#36 from ylee/ylee/manual

Updates to repository README with basic functionality for training and running inference on our custom METRABS models.

README.md

@@ -1,13 +1,47 @@
 # MeTRO Implementation
 1. Setup environment in MLP
 https://roblox.atlassian.net/wiki/spaces/~ylee/pages/2106262747/MeTRo+training+setup+on+MLP
-
-2. Visualization
+    * The datasets are organized within /home/jovyan/data/metrabs-processed/
+    * The trained models are stored in /home/jovyan/runs/metrabs-exp/$log_dir/model/
+
+2. Training Examples
+    1. Trained models are listed in https://docs.google.com/spreadsheets/d/1s_EU91shYbMo9f5R5C5338an3WXiY3TREv5LAt8Gn-0/edit#gid=1046801698
+    2. As of 08/14/2023, we have chosen three models that offer the best trade-off, and you can find their respective training commands by `grep command ${ckpt_dir}/log.txt`.
+        1. Large model(116M): `./main.py --train --dataset=sway --train-on=trainval --seed=1 --training-steps=1000000 --model-class=Metrabs --workers=30 --validate-period=50000 --final-transposed-conv=1 --backbone=mobilenetV3Small --logdir=sway_ncnn/256in_ms_w1_up_pv05_abs0_scratch_2d --proc-side=256 --output-upper-joints --upper-bbox --partial-visibility-prob=0.5 --mobilenet-alpha=1.0 --zero-padding=0 --no-geom-aug --occlude-aug-prob=0 --occlude-aug-prob-2d=0 --dtype=float32 --absloss-factor=0 --data-format=NCHW --input-float32 --init=scratch`
+        2. Medium model(46M): `./main.py --train --dataset=sway --train-on=trainval --seed=1 --training-steps=1000000 --model-class=Metrabs --workers=30 --validate-period=50000 --final-transposed-conv=1 --backbone=mobilenetV3Small --logdir=sway_ncnn/160in_ms_w1_up_pv05_abs0_scratch --proc-side=160 --output-upper-joints --upper-bbox --partial-visibility-prob=0.5 --mobilenet-alpha=1.0 --zero-padding=0 --occlude-aug-prob-2d=0 --dtype=float32 --absloss-factor=0 --data-format=NCHW --input-float32 --init=scratch`
+        3. Small model(20M): `./main.py --train --dataset=sway --train-on=trainval --seed=1 --training-steps=1000000 --model-class=Metrabs --workers=30 --validate-period=50000 --final-transposed-conv=1 --backbone=mobilenetV3Small --logdir=sway_ncnn/112in_ms_w075_up_pv05_abs0_scratch --proc-side=112 --output-upper-joints --upper-bbox --partial-visibility-prob=0.5 --mobilenet-alpha=0.75 --zero-padding=0 --occlude-aug-prob-2d=0 --dtype=float32 --absloss-factor=0 --data-format=NCHW --input-float32 --init=scratch`
+    3. You can manually export a checkpoint into pb format by `ckptno=20001; ./main.py --export-file=model$ckptno --dataset=sway --checkpoint-dir=sway_ncnn/256in_ms_w1_up_pv05_abs0_ckpt --load-path=ckpt-$ckptno --model-class=Metrabs --final-transposed-conv=1 --init=scratch --backbone=mobilenetV3Small --proc-side=256 --mobilenet-alpha=1.0 --output-upper-joints --data-format=NCHW --input-float32 --dtype=float32`
+        * Please ensure that the `data format`, `input-float32`, and `dtype` are consistent with the training settings when exporting the model.
+
+3. Visualization
 `python src/viz.py model1 .. modelN h36m|input_path output_path`
 
-3. Training Examples
-
 4. Evaluation
+    1. Inference results are stored in the model folder under the filename `predictions_sway.npz` by `metrabs/src$ ./main.py --predict --dataset=sway --checkpoint-dir=sway_ncnn/160in_ms_w1_up_pv05_abs0_scratch --backbone=mobilenetV3Small --init=scratch --mobilenet-alpha=1.0 --model-class=Metrabs --proc-side=160 --output-upper-joints --upper-bbox --upper-bbox-ratio 0.5 0.5 --data-format=NCHW --input-float32`
+    2. MPJPE is calculated by `metrabs/src$ python -m eval_scripts.eval_sway --pred-path=/home/jovyan/runs/metrabs-exp/sway_ncnn/160in_ms_w1_up_pv05_abs0_scratch/predictions_sway.npz --root-last`
+
+5. Data preparation
+You can incorporate 3D datasets by following these steps:
+    1. Integrate a offline dataprocessing function into the `data/datasets3d.py` file. Refer to the example at https://github.rbx.com/ylee/metrabs/pull/10 for guidance
+    2. Include the offline data processing function within the `data/sway.py` file.
+    3. To evaluate models using the new data, develop the evaluation script within the `eval_scripts/eval_sway.py` file.
+
+6. TF Model conversion to ONNX and NCNN
+    1. Convert TF model to ONNX model
+    `PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python python3 -m tf2onnx.convert --saved-model model_path --output output_path --opset 15`    
+    example: `PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python python3 -m tf2onnx.convert --saved-model ./runs/metrabs-exp/sway_ncnn/256in_ms_w1_up_pv05_abs0_scratch_2d/model/ --output ./output/256in_ms_w1_up_pv05_abs0_scratch_2d.onnx --opset 15`
+
+    2. Dynamic input to fixed input
+    `python3 -m onnxruntime.tools.make_dynamic_shape_fixed input_onnx_path output_onnx_path --input_name input_2 --input_shape N,C,H,W`    
+    example: `python3 -m onnxruntime.tools.make_dynamic_shape_fixed output/256in_ms_w1_up_pv05_abs0_scratch_2d.onnx 256in_ms_w1_up_pv05_abs0_scratch_2d.onnx --input_name input_2 --input_shape 1,3,256,256`
+
+    3. Simplify ONNX model
+    https://convertmodel.com/#input=onnx&output=onnx
+    Visit this link and choose the input format :onnx , and the out format: onnx. Check all the options and click the convert button. You can download simplified onnx model.
+
+    4. Convert simplified ONNX model to NCNN model
+    https://convertmodel.com/#input=onnx&output=onnx
+    Choose output format : ncnn, and the input format : onnx. Uncheck all the options and click the convert button. You can download bin and param files. 
 
 
 # MeTRAbs Absolute 3D Human Pose Estimator