From 3e4c3670b10d07d8fdbaab982afc469749cb0de6 Mon Sep 17 00:00:00 2001
From: yinwanqi <yinwanqi@sensetime.com>
Date: Thu, 20 Jul 2023 17:58:46 +0800
Subject: [PATCH] add docs for inference

---
 README.md               | 29 ++++++++++++++++++++++++++++-
 main/inference.py       |  7 +++----
 main/slurm_inference.sh | 11 +++++------
 3 files changed, 36 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 298f18a..ac52f41 100644
--- a/README.md
+++ b/README.md
@@ -92,7 +92,18 @@ SMPLer-X/
 │               └──SMPLX_FEMALE.npz
 ├── data/
 ├── main/
-├── pretrained_models/  # pretrained ViT-Pose models
+├── demo/  
+│   ├── videos/       
+│   ├── images/      
+│   └── results/ 
+├── pretrained_models/  # pretrained ViT-Pose, SMPLer_X and mmdet models
+│   ├── mmdet/
+│   │   ├──faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth
+│   │   └──mmdet_faster_rcnn_r50_fpn_coco.py
+│   ├── smpler_x_s32.pth.tar
+│   ├── smpler_x_b32.pth.tar
+│   ├── smpler_x_l32.pth.tar
+│   ├── smpler_x_h32.pth.tar
 │   ├── vitpose_small.pth
 │   ├── vitpose_base.pth
 │   ├── vitpose_large.pth
@@ -131,6 +142,21 @@ SMPLer-X/
     ├── UP3D/
     └── preprocessed_datasets/  # HumanData files
 ```
+## Inference 
+- Place the video to be inferenced under ROOT/demo/videos
+- Prepare the pretrained models to be used for inference under ROOT/pretrained_models 
+- Prepare the mmdet pretrained model and config under ROOT/pretrained_models 
+- Inference out put will be placed in ROOT/demo/results
+
+```bash
+cd main
+sh slurm_inference.sh {VIDEO_FILE} {FORMAT} {FPS} {PRETRAINED_CKPT} 
+
+# For inferencing test_video.mp4 (24FPS) with smpler_x_h32
+sh slurm_inference.sh test_video mp4 24 smpler_x_h32
+
+```
+
 
 ## Training
 ```bash
@@ -155,6 +181,7 @@ sh slurm_test.sh {JOB_NAME} {NUM_GPU} {TRAIN_OUTPUT_DIR} {CKPT_ID}
 - NUM_GPU = 1 is recommended for testing
 - Logs and results  will be saved to `../output/test_{JOB_NAME}_ep{CKPT_ID}_{TEST_DATSET}`
 
+
 ## References
 - [Hand4Whole](https://github.com/mks0601/Hand4Whole_RELEASE)
 - [OSX](https://github.com/IDEA-Research/OSX)
diff --git a/main/inference.py b/main/inference.py
index 3c68d1e..87f888c 100644
--- a/main/inference.py
+++ b/main/inference.py
@@ -21,8 +21,7 @@ def parse_args():
     parser = argparse.ArgumentParser()
     parser.add_argument('--num_gpus', type=int, dest='num_gpus')
     parser.add_argument('--exp_name', type=str, default='output/test')
-    parser.add_argument('--result_path', type=str, default='output/test')
-    parser.add_argument('--ckpt_idx', type=int, default=0)
+    parser.add_argument('--pretrained_model', type=int, default=0)
     parser.add_argument('--testset', type=str, default='EHF')
     parser.add_argument('--agora_benchmark', type=str, default='na')
     parser.add_argument('--img_path', type=str, default='input.png')
@@ -43,8 +42,8 @@ def parse_args():
 def main():
 
     args = parse_args()
-    config_path = osp.join('../output',args.result_path, 'code', 'config_base.py')
-    ckpt_path = osp.join('../output', args.result_path, 'model_dump', f'snapshot_{int(args.ckpt_idx)}.pth.tar')
+    config_path = osp.join('./config', f'config_{args.pretrained_model}.py')
+    ckpt_path = osp.join('../pretrained_models', f'{args.pretrained_model}.pth.tar')
 
     cfg.get_config_fromfile(config_path)
     cfg.update_test_config(args.testset, args.agora_benchmark, shapy_eval_split=None, 
diff --git a/main/slurm_inference.sh b/main/slurm_inference.sh
index a010ba6..7fa2544 100644
--- a/main/slurm_inference.sh
+++ b/main/slurm_inference.sh
@@ -4,10 +4,10 @@ set -x
 PARTITION=Zoetrope
 
 INPUT_VIDEO=$1
-APPENDIX=$2
+FORMAT=$2
 FPS=$3
-RES_PATH=$4
-CKPT=$5
+CKPT=$4
+
 GPUS=1
 JOB_NAME=inference_${INPUT_VIDEO}
 
@@ -21,7 +21,7 @@ SAVE_DIR=../demo/results/${INPUT_VIDEO}
 # video to images
 mkdir $IMG_PATH
 mkdir $SAVE_DIR
-ffmpeg -i ../demo/videos/${INPUT_VIDEO}.${APPENDIX} -f image2 -vf fps=${FPS}/1 -qscale 0 ../demo/images/${INPUT_VIDEO}/%06d.jpg 
+ffmpeg -i ../demo/videos/${INPUT_VIDEO}.${FORMAT} -f image2 -vf fps=${FPS}/1 -qscale 0 ../demo/images/${INPUT_VIDEO}/%06d.jpg 
 
 end_count=$(find "$IMG_PATH" -type f | wc -l)
 echo $end_count
@@ -39,8 +39,7 @@ srun -p ${PARTITION} \
     python inference.py \
     --num_gpus ${GPUS_PER_NODE} \
     --exp_name output/demo_${JOB_NAME} \
-    --result_path ${RES_PATH} \
-    --ckpt_idx ${CKPT} \
+    --pretrained_model ${CKPT} \
     --agora_benchmark agora_model \
     --img_path ${IMG_PATH} \
     --start 1 \