File size: 1,757 Bytes
20b4034 | 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 | #!/usr/bin/env bash
set -euo pipefail
UMM_ROOT="${UMM_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
SOURCE_IMAGES="${SOURCE_IMAGES:-/rczhang/rczhang/zhangrch/video_gen/gen_dataset/covt/images}"
FILES_PER_SHARD="${FILES_PER_SHARD:-12000}"
LIST_DIR="${UMM_ROOT}/data/covt/shard_lists"
SHARD_DIR="${UMM_ROOT}/data/covt/shards"
unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY all_proxy ALL_PROXY
mkdir -p "${LIST_DIR}" "${SHARD_DIR}"
if [[ ! -d "${SOURCE_IMAGES}" ]]; then
echo "Missing source image directory: ${SOURCE_IMAGES}" >&2
exit 1
fi
if ! compgen -G "${LIST_DIR}/images-*.txt" >/dev/null; then
find "${SOURCE_IMAGES}" -maxdepth 1 -type f -printf '%f\n' \
| LC_ALL=C sort \
| split -d -a 5 -l "${FILES_PER_SHARD}" --additional-suffix=.txt - "${LIST_DIR}/images-"
fi
for list_path in "${LIST_DIR}"/images-*.txt; do
shard_name="$(basename "${list_path}" .txt)"
shard_path="${SHARD_DIR}/${shard_name}.tar"
checksum_path="${shard_path}.sha256"
temporary_path="${shard_path}.partial"
if [[ -s "${shard_path}" && -s "${checksum_path}" ]]; then
echo "reuse ${shard_name}"
continue
fi
echo "build ${shard_name}"
tar -C "${SOURCE_IMAGES}" -cf "${temporary_path}" -T "${list_path}"
mv "${temporary_path}" "${shard_path}"
(
cd "${SHARD_DIR}"
sha256sum "$(basename "${shard_path}")" > "$(basename "${checksum_path}")"
)
done
expected="$(find "${SOURCE_IMAGES}" -maxdepth 1 -type f | wc -l)"
listed="$(awk 'END {print NR}' "${LIST_DIR}"/images-*.txt)"
if [[ "${expected}" != "${listed}" ]]; then
echo "Image list mismatch: source=${expected}, listed=${listed}" >&2
exit 1
fi
echo "Built $(find "${SHARD_DIR}" -maxdepth 1 -name 'images-*.tar' | wc -l) image shards for ${listed} images"
|