3RScan/run_3rscan_sonata.sh

#!/usr/bin/env bash
set -euo pipefail

# 3RScan .pth → NPZ (coord+color+normal) → Sonata feat
# Input:  scan_data/pcd_with_global_alignment/*.pth
# Output: sonata_feat/<scan_id>_down.npz

export CUDA_VISIBLE_DEVICES=0
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128

PYTHON="/home/lws0111/Sonata_run/venv/bin/python"
FEAT_SCRIPT="/home/lws0111/Sonata_run/sonata/code/extract_features.py"

BASE="/home/lws0111/Warehouse/3RScan"
PTH_DIR="${BASE}/scan_data/pcd_with_global_alignment"
CONVERTED_DIR="${BASE}/converted"
FEAT_DIR="${BASE}/sonata_feat"

mkdir -p "${CONVERTED_DIR}" "${FEAT_DIR}"

# ---- Step 1: .pth → .npz (coord + color + normal) ----
echo "========================================"
echo "[Step 1] Convert .pth → .npz (with normals)"
echo "========================================"

${PYTHON} -u -c "
import os, numpy as np, torch, open3d as o3d

pth_dir = '${PTH_DIR}'
out_dir = '${CONVERTED_DIR}'

files = sorted(f for f in os.listdir(pth_dir) if f.endswith('.pth'))
total = len(files)
done, skip = 0, 0

for f in files:
    scan_id = f.replace('.pth', '')
    out_path = os.path.join(out_dir, scan_id + '.npz')
    if os.path.exists(out_path):
        skip += 1
        continue

    data = torch.load(os.path.join(pth_dir, f), map_location='cpu', weights_only=False)
    coord = np.asarray(data[0], dtype=np.float32)
    color = np.asarray(data[1], dtype=np.float32)
    coord = np.nan_to_num(coord, nan=0.0, posinf=0.0, neginf=0.0)
    color = np.nan_to_num(color, nan=0.0, posinf=0.0, neginf=0.0)
    if color.max() > 1.1:
        color = color / 255.0
    color = np.clip(color, 0.0, 1.0)

    # Estimate normals with Open3D
    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coord)
    pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=0.05, max_nn=30))
    normal = np.asarray(pcd.normals, dtype=np.float32)

    np.savez_compressed(out_path, coord=coord, color=color, normal=normal)
    done += 1
    if done % 20 == 0 or done == 1:
        print(f'  [{skip+done}/{total}] {scan_id} ({coord.shape[0]} pts)', flush=True)

print(f'[Step 1 Done] converted={done}, skipped={skip}')
"

# ---- Step 2: Sonata feature extraction ----
echo ""
echo "========================================"
echo "[Step 2] Sonata feature extraction"
echo "========================================"

mapfile -t ALL_NPZ < <(ls "${CONVERTED_DIR}"/*.npz 2>/dev/null | sort)
total=${#ALL_NPZ[@]}

todo=0
skip=0
for npz in "${ALL_NPZ[@]}"; do
  scan_id="$(basename "${npz}" .npz)"
  feat="${FEAT_DIR}/${scan_id}_down.npz"
  if [ -f "$feat" ] && [ -s "$feat" ]; then
    ((skip+=1))
  else
    ((todo+=1))
  fi
done
echo "[I] Total=${total}, todo=${todo}, already_done=${skip}"

[ "$todo" -eq 0 ] && echo "[DONE] All done!" && exit 0

processed=0
failed=0
for npz in "${ALL_NPZ[@]}"; do
  scan_id="$(basename "${npz}" .npz)"
  feat="${FEAT_DIR}/${scan_id}_down.npz"

  [ -f "$feat" ] && [ -s "$feat" ] && continue

  ((processed+=1))
  echo "[${processed}/${todo}] ${scan_id}"

  # Try 1: no-flash + patch=512 + fp16
  set +e
  ${PYTHON} "${FEAT_SCRIPT}" \
    --input "${npz}" \
    --output "${feat}" \
    --fp16 --no-flash --patch 512 2>&1 | tail -3
  rc=$?
  set -e

  if [ $rc -eq 0 ] && [ -f "$feat" ] && [ -s "$feat" ]; then
    echo "  [OK]"
    continue
  fi

  # Try 2: patch=256
  echo "  [RETRY] patch=256 ..."
  rm -f "${feat}" 2>/dev/null || true
  set +e
  ${PYTHON} "${FEAT_SCRIPT}" \
    --input "${npz}" \
    --output "${feat}" \
    --fp16 --no-flash --patch 256 2>&1 | tail -3
  rc=$?
  set -e

  if [ $rc -eq 0 ] && [ -f "$feat" ] && [ -s "$feat" ]; then
    echo "  [OK] (patch=256)"
    continue
  fi

  echo "  [FAIL] ${scan_id}"
  rm -f "${feat}" 2>/dev/null || true
  ((failed+=1))
done

echo ""
echo "[DONE] processed=${processed}, failed=${failed}"
echo "  Output: ${FEAT_DIR}/"