diff options
| -rwxr-xr-x | scripts/train_cnn_v2_full.sh | 111 |
1 files changed, 81 insertions, 30 deletions
diff --git a/scripts/train_cnn_v2_full.sh b/scripts/train_cnn_v2_full.sh index 7e5aeee..b3ece8a 100755 --- a/scripts/train_cnn_v2_full.sh +++ b/scripts/train_cnn_v2_full.sh @@ -1,12 +1,24 @@ #!/bin/bash # Complete CNN v2 Training Pipeline # Train → Export → Build → Validate +# Usage: ./train_cnn_v2_full.sh [--validate [checkpoint.pth]] set -e PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" cd "$PROJECT_ROOT" +# Parse arguments +VALIDATE_ONLY=false +VALIDATE_CHECKPOINT="" + +if [ "$1" = "--validate" ]; then + VALIDATE_ONLY=true + if [ -n "$2" ]; then + VALIDATE_CHECKPOINT="$2" + fi +fi + # Configuration INPUT_DIR="training/input" TARGET_DIR="training/target_2" @@ -28,15 +40,22 @@ DETECTOR="harris" KERNEL_SIZES="3 3 3" CHANNELS="8 4 4" -echo "=== CNN v2 Complete Training Pipeline ===" -echo "Input: $INPUT_DIR" -echo "Target: $TARGET_DIR" -echo "Epochs: $EPOCHS" -echo "Checkpoint interval: $CHECKPOINT_EVERY" -echo "" +if [ "$VALIDATE_ONLY" = true ]; then + echo "=== CNN v2 Validation Only ===" + echo "Skipping training, using existing weights" + echo "" +else + echo "=== CNN v2 Complete Training Pipeline ===" + echo "Input: $INPUT_DIR" + echo "Target: $TARGET_DIR" + echo "Epochs: $EPOCHS" + echo "Checkpoint interval: $CHECKPOINT_EVERY" + echo "" +fi -# Step 1: Train model -echo "[1/4] Training CNN v2 model..." +if [ "$VALIDATE_ONLY" = false ]; then + # Step 1: Train model + echo "[1/4] Training CNN v2 model..." python3 training/train_cnn_v2.py \ --input "$INPUT_DIR" \ --target "$TARGET_DIR" \ @@ -79,30 +98,51 @@ if [ $? -ne 0 ]; then fi echo "" +fi # End of training/export section + +# Determine which checkpoint to use +if [ "$VALIDATE_ONLY" = true ]; then + if [ -n "$VALIDATE_CHECKPOINT" ]; then + FINAL_CHECKPOINT="$VALIDATE_CHECKPOINT" + else + # Use latest checkpoint + FINAL_CHECKPOINT=$(ls -t "$CHECKPOINT_DIR"/checkpoint_epoch_*.pth | head -1) + fi + echo "Using checkpoint: $FINAL_CHECKPOINT" + echo "" +fi # Step 3: Rebuild with new shaders -echo "[3/4] Rebuilding demo with new shaders..." -cmake --build build -j4 --target demo64k > /dev/null 2>&1 +if [ "$VALIDATE_ONLY" = false ]; then + echo "[3/4] Rebuilding demo with new shaders..." + cmake --build build -j4 --target demo64k > /dev/null 2>&1 -if [ $? -ne 0 ]; then - echo "Error: Build failed" - exit 1 -fi + if [ $? -ne 0 ]; then + echo "Error: Build failed" + exit 1 + fi -echo " → Build complete" -echo "" + echo " → Build complete" + echo "" +fi # Step 4: Visual assessment - process final checkpoint only -echo "[4/4] Visual assessment on all input images..." -mkdir -p "$VALIDATION_DIR" +if [ "$VALIDATE_ONLY" = true ]; then + echo "Validation on all input images (using existing weights)..." +else + echo "[4/4] Visual assessment on all input images..." +fi -echo " Processing final checkpoint: $FINAL_CHECKPOINT" +mkdir -p "$VALIDATION_DIR" +echo " Using checkpoint: $FINAL_CHECKPOINT" -# Export final checkpoint shaders -python3 training/export_cnn_v2_weights.py "$FINAL_CHECKPOINT" \ - --output-weights workspaces/main/cnn_v2_weights.bin > /dev/null 2>&1 +# Export weights only if not in validate mode +if [ "$VALIDATE_ONLY" = false ]; then + python3 training/export_cnn_v2_weights.py "$FINAL_CHECKPOINT" \ + --output-weights workspaces/main/cnn_v2_weights.bin > /dev/null 2>&1 +fi -# Rebuild with final weights +# Build cnn_test cmake --build build -j4 --target cnn_test > /dev/null 2>&1 # Process all input images @@ -112,19 +152,30 @@ for input_image in "$INPUT_DIR"/*.png; do build/cnn_test "$input_image" "$VALIDATION_DIR/${basename}_output.png" 2>/dev/null done -cmake --build build -j4 --target demo64k > /dev/null 2>&1 +# Build demo only if not in validate mode +if [ "$VALIDATE_ONLY" = false ]; then + cmake --build build -j4 --target demo64k > /dev/null 2>&1 +fi echo "" -echo "=== Training Pipeline Complete ===" +if [ "$VALIDATE_ONLY" = true ]; then + echo "=== Validation Complete ===" +else + echo "=== Training Pipeline Complete ===" +fi echo "" echo "Results:" -echo " - Checkpoints: $CHECKPOINT_DIR" +if [ "$VALIDATE_ONLY" = false ]; then + echo " - Checkpoints: $CHECKPOINT_DIR" + echo " - Final weights: workspaces/main/cnn_v2_weights.bin" +fi echo " - Validation outputs: $VALIDATION_DIR" -echo " - Final weights: workspaces/main/cnn_v2_weights.bin" echo "" echo "Opening results directory..." open "$VALIDATION_DIR" 2>/dev/null || xdg-open "$VALIDATION_DIR" 2>/dev/null || true -echo "" -echo "Run demo to see final result:" -echo " ./build/demo64k" +if [ "$VALIDATE_ONLY" = false ]; then + echo "" + echo "Run demo to see final result:" + echo " ./build/demo64k" +fi |