#!/bin/bash
# Provenance record: the exact pipeline that produced the TLT12 release.
#
# This file documents how the shards in this release were made. It is a record,
# not a turnkey script -- paths are cluster-specific and rendering the full set
# takes several GPU-days. See `provenance.json` for the generator commit.
#
# Generator: https://github.com/fzenke/3D_latent_traversal_dataset_generator
# Requires:  BlenderProc + a local copy of ShapeNetCore.v2 (https://shapenet.org)
set -euo pipefail

SHAPENET_PATH=/path/to/ShapeNetCoreV2
OUTPUT_DIR=/path/to/tlt12
WDS_DIR=/path/to/tlt12_wds

# ── 1. Select the objects: 12 synsets, all objects in each ───────────────────
ARGS=""
for SYNSETID in 02691156 02958343 02954340 03001627 03261776 03467517 \
                04530566 03642806 03790512 03797390 03928116 04401088; do
  ARGS="$ARGS --synset $SYNSETID all"
done
python sample_objects.py \
  --models-path "$SHAPENET_PATH" \
  --output objects.npy \
  $ARGS

# ── 2. Keep 10 objects per synset and split across 60 SLURM array tasks ──────
python split_objects.py \
  --objects objects.npy \
  --n-jobs 60 \
  --output-dir splits \
  --max 10

# ── 3. Render ────────────────────────────────────────────────────────────────
# 128x128 px, 32 frames per sequence, 1000 sequences per object.
# One array task per split file; see submit_array.sh in the generator repo.
#
#   sbatch --array=0-59 submit_array.sh
#
# Each task runs, with SLURM_ARRAY_TASK_ID selecting the split:
#
#   blenderproc run generate_traversals.py \
#     --models-path "$SHAPENET_PATH" \
#     --output-dir  "$OUTPUT_DIR" \
#     --objects     splits/objects_NNN.npy \
#     --image-size 128 \
#     --n-frames 32 \
#     --seqs-per-object 1000 \
#     --multi-factor \
#     --seed NNN
#
# Face deduplication (--dedupe-faces) is on by default and matters: most
# ShapeNetCore.v2 meshes store every face twice on the same vertices, which
# produces flickering z-fighting patches as the object rotates. See the
# "Known issues and differences from 3DIEBench" section of ../README.md.

# ── 4. Merge the per-task metadata ───────────────────────────────────────────
python merge_metadata.py -v --output-dir "$OUTPUT_DIR"

# ── 5. Pack into WebDataset shards with a 80/10/10 split ─────────────────────
python package_wds.py \
  --dataset-dir "$OUTPUT_DIR" \
  --output-dir  "$WDS_DIR" \
  --val-fraction 0.1 \
  --test-fraction 0.1 \
  --shard-maxcount 1000