Spaces:
Running
Running
File size: 2,769 Bytes
ba23d94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 | #!/bin/bash
# Run surface-normal estimation on a directory of images.
cd "$(dirname "$(realpath "$0")")/../.." || exit
SAPIENS_CHECKPOINT_ROOT="${SAPIENS_CHECKPOINT_ROOT:-${HOME}/sapiens2_host}"
#----------------------------set your input and output directories-------------------------
INPUT='./demo/data/itw_videos/reel1'
OUTPUT="${HOME}/Desktop/sapiens2/normal/Outputs/vis/itw_videos/reel1"
#--------------------------MODEL CARD (uncomment one)---------------------------------------
# MODEL_NAME='sapiens2_0.4b'; CHECKPOINT="${SAPIENS_CHECKPOINT_ROOT}/normal/sapiens2_0.4b_normal.safetensors"
# MODEL_NAME='sapiens2_0.8b'; CHECKPOINT="${SAPIENS_CHECKPOINT_ROOT}/normal/sapiens2_0.8b_normal.safetensors"
MODEL_NAME='sapiens2_1b'; CHECKPOINT="${SAPIENS_CHECKPOINT_ROOT}/normal/sapiens2_1b_normal.safetensors"
# MODEL_NAME='sapiens2_5b'; CHECKPOINT="${SAPIENS_CHECKPOINT_ROOT}/normal/sapiens2_5b_normal.safetensors"
DATASET='metasim_render_people'
MODEL="${MODEL_NAME}_normal_${DATASET}-1024x768"
CONFIG_FILE="configs/normal/${DATASET}/${MODEL}.py"
OUTPUT="${OUTPUT}/${MODEL_NAME}"
##-------------------------------------inference--------------------------------------------
RUN_FILE='tools/vis/vis_normal.py'
JOBS_PER_GPU=3; GPU_IDS=(0 1 2 3 4 5 6 7)
# JOBS_PER_GPU=1; GPU_IDS=(0)
TOTAL_JOBS=$((JOBS_PER_GPU * ${#GPU_IDS[@]}))
IMAGE_LIST="${INPUT}/image_list.txt"
find "${INPUT}" -type f \( -iname \*.jpg -o -iname \*.jpeg -o -iname \*.png \) | sort > "${IMAGE_LIST}"
if [ ! -s "${IMAGE_LIST}" ]; then
echo "No images found at ${INPUT}"
exit 1
fi
NUM_IMAGES=$(wc -l < "${IMAGE_LIST}")
IMAGES_PER_FILE=$((NUM_IMAGES / TOTAL_JOBS))
EXTRA_IMAGES=$((NUM_IMAGES % TOTAL_JOBS))
export TF_CPP_MIN_LOG_LEVEL=2
echo "Distributing ${NUM_IMAGES} image paths into ${TOTAL_JOBS} jobs."
current_line=1
for ((i=0; i<TOTAL_JOBS; i++)); do
TEXT_FILE="${INPUT}/image_paths_$((i+1)).txt"
if [ $i -lt $EXTRA_IMAGES ]; then
images_for_this_job=$((IMAGES_PER_FILE + 1))
else
images_for_this_job=$IMAGES_PER_FILE
fi
if [ $images_for_this_job -gt 0 ]; then
sed -n "${current_line},$((current_line + images_for_this_job - 1))p" "${IMAGE_LIST}" > "${TEXT_FILE}"
current_line=$((current_line + images_for_this_job))
else
touch "${TEXT_FILE}"
fi
done
for ((i=0; i<TOTAL_JOBS; i++)); do
GPU_ID=${GPU_IDS[$((i % ${#GPU_IDS[@]}))]}
CMD="CUDA_VISIBLE_DEVICES=${GPU_ID} python ${RUN_FILE} \
${CONFIG_FILE} \
${CHECKPOINT} \
--save_pred \
--input \"${INPUT}/image_paths_$((i+1)).txt\" \
--output \"${OUTPUT}\""
[ "$TOTAL_JOBS" -gt 1 ] && CMD="$CMD &"
eval $CMD
sleep 1
done
wait
rm "${IMAGE_LIST}"
for ((i=0; i<TOTAL_JOBS; i++)); do
rm "${INPUT}/image_paths_$((i+1)).txt"
done
echo "Output directory: ${OUTPUT}"
|