blob: fc9355a547a3ff07943f6a427600a825e10eb160 (
plain)
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
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
|
#!/bin/bash
# Complete CNN v2 Training Pipeline
# Train → Export → Build → Validate
# Usage: ./train_cnn_v2_full.sh [OPTIONS]
#
# OPTIONS:
# (none) Run complete pipeline: train → export → build → validate
# --validate Validate only (skip training, use existing weights)
# --validate CHECKPOINT Validate with specific checkpoint file
# --help Show this help message
#
# Examples:
# ./train_cnn_v2_full.sh
# ./train_cnn_v2_full.sh --validate
# ./train_cnn_v2_full.sh --validate checkpoints/checkpoint_epoch_50.pth
set -e
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$PROJECT_ROOT"
# Parse arguments
VALIDATE_ONLY=false
VALIDATE_CHECKPOINT=""
if [ "$1" = "--help" ] || [ "$1" = "-h" ]; then
head -20 "$0" | grep "^#" | grep -v "^#!/" | sed 's/^# *//'
exit 0
fi
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"
CHECKPOINT_DIR="checkpoints"
VALIDATION_DIR="validation_results"
EPOCHS=100
CHECKPOINT_EVERY=5
BATCH_SIZE=16
# Patch-based training (default)
PATCH_SIZE=32
PATCHES_PER_IMAGE=64
DETECTOR="harris"
# Full-image training (alternative - uncomment to use)
# FULL_IMAGE="--full-image"
# IMAGE_SIZE=256
KERNEL_SIZES="3 3 3"
CHANNELS="8 4 4"
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
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" \
--patch-size $PATCH_SIZE \
--patches-per-image $PATCHES_PER_IMAGE \
--detector $DETECTOR \
--kernel-sizes $KERNEL_SIZES \
--channels $CHANNELS \
--epochs $EPOCHS \
--batch-size $BATCH_SIZE \
--checkpoint-dir "$CHECKPOINT_DIR" \
--checkpoint-every $CHECKPOINT_EVERY \
$FULL_IMAGE
if [ $? -ne 0 ]; then
echo "Error: Training failed"
exit 1
fi
echo ""
echo "Training complete!"
echo ""
# Step 2: Export final checkpoint to shaders
FINAL_CHECKPOINT="$CHECKPOINT_DIR/checkpoint_epoch_${EPOCHS}.pth"
if [ ! -f "$FINAL_CHECKPOINT" ]; then
echo "Warning: Final checkpoint not found, using latest available..."
FINAL_CHECKPOINT=$(ls -t "$CHECKPOINT_DIR"/checkpoint_epoch_*.pth | head -1)
fi
echo "[2/4] Exporting final checkpoint to WGSL shaders..."
echo "Checkpoint: $FINAL_CHECKPOINT"
python3 training/export_cnn_v2_shader.py "$FINAL_CHECKPOINT" \
--output-dir workspaces/main/shaders
if [ $? -ne 0 ]; then
echo "Error: Shader export failed"
exit 1
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
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
echo " → Build complete"
echo ""
fi
# Step 4: Visual assessment - process final checkpoint only
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
mkdir -p "$VALIDATION_DIR"
echo " Using checkpoint: $FINAL_CHECKPOINT"
# 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
# Build cnn_test
cmake --build build -j4 --target cnn_test > /dev/null 2>&1
# Process all input images
for input_image in "$INPUT_DIR"/*.png; do
basename=$(basename "$input_image" .png)
echo " Processing $basename..."
build/cnn_test "$input_image" "$VALIDATION_DIR/${basename}_output.png" 2>/dev/null
done
# 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 ""
if [ "$VALIDATE_ONLY" = true ]; then
echo "=== Validation Complete ==="
else
echo "=== Training Pipeline Complete ==="
fi
echo ""
echo "Results:"
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 ""
echo "Opening results directory..."
open "$VALIDATION_DIR" 2>/dev/null || xdg-open "$VALIDATION_DIR" 2>/dev/null || true
if [ "$VALIDATE_ONLY" = false ]; then
echo ""
echo "Run demo to see final result:"
echo " ./build/demo64k"
fi
|
