419 lines
14 KiB
Bash
419 lines
14 KiB
Bash
#!/usr/bin/env bash
|
|
# Fine-tune Tesseract 5 LSTM on Fontin / Fontin SmallCaps for POE2
|
|
#
|
|
# IMPORTANT: Run inside WSL Ubuntu (NOT docker-desktop).
|
|
#
|
|
# From PowerShell:
|
|
# wsl -d Ubuntu-22.04 bash /mnt/c/Users/boki/repos/poe2trade/tools/training/train.sh
|
|
#
|
|
# Or from inside WSL Ubuntu:
|
|
# cd /mnt/c/Users/boki/repos/poe2trade/tools/training
|
|
# bash train.sh
|
|
#
|
|
# Prerequisites:
|
|
# - WSL Ubuntu installed: wsl --install -d Ubuntu-22.04
|
|
# - Internet access (downloads fonts + Tesseract tools if needed)
|
|
|
|
set -euo pipefail
|
|
|
|
# ── Sanity checks ────────────────────────────────────────────────────────────
|
|
|
|
if [ -z "${BASH_VERSION:-}" ]; then
|
|
echo "ERROR: This script requires bash. Run with: bash train.sh"
|
|
echo "If you're in docker-desktop WSL, switch to Ubuntu:"
|
|
echo " wsl -d Ubuntu-22.04 bash /mnt/c/Users/boki/repos/poe2trade/tools/training/train.sh"
|
|
exit 1
|
|
fi
|
|
|
|
# Detect which WSL distro we're in
|
|
if grep -qi alpine /etc/os-release 2>/dev/null || [ -f /etc/alpine-release ]; then
|
|
echo "ERROR: You're in docker-desktop (Alpine). Use Ubuntu instead:"
|
|
echo " wsl -d Ubuntu-22.04 bash /mnt/c/Users/boki/repos/poe2trade/tools/training/train.sh"
|
|
exit 1
|
|
fi
|
|
|
|
# ── Resolve paths ────────────────────────────────────────────────────────────
|
|
|
|
# Handle both /mnt/c/ (Ubuntu) and /mnt/host/c/ (docker-desktop) paths
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
|
|
TESSDATA_DIR="$PROJECT_ROOT/tools/OcrDaemon/tessdata"
|
|
TRAINING_TEXT="$SCRIPT_DIR/poe2_training_text.txt"
|
|
WORK_DIR="$HOME/poe2-tesseract-training"
|
|
FONT_DIR="/usr/local/share/fonts/fontin"
|
|
MAX_ITERATIONS=800
|
|
TARGET_ERROR=0.01
|
|
|
|
if [ ! -f "$TRAINING_TEXT" ]; then
|
|
echo "ERROR: Training text not found at: $TRAINING_TEXT"
|
|
exit 1
|
|
fi
|
|
|
|
if [ ! -f "$TESSDATA_DIR/eng.traineddata" ]; then
|
|
echo "ERROR: eng.traineddata not found at: $TESSDATA_DIR/"
|
|
exit 1
|
|
fi
|
|
|
|
echo "=== POE2 Tesseract Fine-Tuning ==="
|
|
echo "Script dir: $SCRIPT_DIR"
|
|
echo "Work dir: $WORK_DIR"
|
|
echo "Training text: $TRAINING_TEXT"
|
|
echo "Tessdata: $TESSDATA_DIR"
|
|
echo ""
|
|
|
|
# ── Step 1: Install Tesseract training tools ─────────────────────────────────
|
|
|
|
install_tesseract_tools() {
|
|
echo "── Step 1: Installing Tesseract training tools ──"
|
|
|
|
if command -v text2image >/dev/null 2>&1 && command -v lstmtraining >/dev/null 2>&1; then
|
|
echo "Tesseract training tools already installed."
|
|
tesseract --version 2>&1 | head -1
|
|
return 0
|
|
fi
|
|
|
|
echo "Installing Tesseract 5.x and training tools..."
|
|
|
|
# Update package list first
|
|
sudo apt-get update -qq
|
|
|
|
# Try PPA for latest Tesseract
|
|
if sudo add-apt-repository -y ppa:alex-p/tesseract-ocr-devel 2>/dev/null; then
|
|
sudo apt-get update -qq
|
|
fi
|
|
|
|
sudo apt-get install -y tesseract-ocr tesseract-ocr-eng \
|
|
libicu-dev libpango1.0-dev libcairo2-dev \
|
|
wget unzip fontconfig 2>&1 | tail -3
|
|
|
|
# Check if we got text2image from the package
|
|
if command -v text2image >/dev/null 2>&1; then
|
|
echo "Installed via apt."
|
|
tesseract --version 2>&1 | head -1
|
|
return 0
|
|
fi
|
|
|
|
# Build from source as fallback
|
|
echo "Package didn't include training tools, building from source..."
|
|
sudo apt-get install -y automake g++ git libtool libleptonica-dev \
|
|
make pkg-config libicu-dev libpango1.0-dev libcairo2-dev 2>&1 | tail -3
|
|
|
|
local BUILD_DIR="$HOME/tesseract-build"
|
|
if [ ! -d "$BUILD_DIR" ]; then
|
|
git clone --depth 1 --branch 5.3.4 https://github.com/tesseract-ocr/tesseract.git "$BUILD_DIR"
|
|
fi
|
|
|
|
pushd "$BUILD_DIR" > /dev/null
|
|
./autogen.sh
|
|
./configure
|
|
make -j"$(nproc)"
|
|
make training
|
|
sudo make install
|
|
sudo make training-install
|
|
sudo ldconfig
|
|
popd > /dev/null
|
|
|
|
echo "Built from source."
|
|
tesseract --version 2>&1 | head -1
|
|
}
|
|
|
|
# ── Step 2: Download and install Fontin fonts ────────────────────────────────
|
|
|
|
install_fonts() {
|
|
echo ""
|
|
echo "── Step 2: Installing Fontin fonts ──"
|
|
|
|
if fc-list 2>/dev/null | grep -qi fontin; then
|
|
echo "Fontin fonts already installed:"
|
|
fc-list | grep -i fontin
|
|
return 0
|
|
fi
|
|
|
|
local FONT_TMP="$HOME/fontin-download"
|
|
mkdir -p "$FONT_TMP"
|
|
pushd "$FONT_TMP" > /dev/null
|
|
|
|
echo "Downloading Fontin fonts..."
|
|
# wfonts.com zip includes Regular, Bold, Italic, SmallCaps (all in one)
|
|
curl -sL "https://static.wfonts.com/download/data/2015/03/10/fontin/fontin.zip" -o fontin.zip || {
|
|
echo "ERROR: Failed to download Fontin fonts."
|
|
echo "Download manually and place .otf files in $FONT_DIR"
|
|
popd > /dev/null
|
|
return 1
|
|
}
|
|
|
|
# Verify it's a real zip
|
|
if ! file fontin.zip | grep -q "Zip archive"; then
|
|
echo "ERROR: Downloaded file is not a zip. Font site may be blocking downloads."
|
|
echo "Download manually from https://www.wfonts.com/font/fontin"
|
|
echo "Place .otf files in $FONT_DIR"
|
|
popd > /dev/null
|
|
return 1
|
|
fi
|
|
|
|
unzip -qo fontin.zip -d fontin/ 2>/dev/null || true
|
|
|
|
sudo mkdir -p "$FONT_DIR"
|
|
# Copy OTF files
|
|
find fontin/ \( -name '*.otf' -o -name '*.OTF' \) \
|
|
-exec sudo cp {} "$FONT_DIR/" \; 2>/dev/null || true
|
|
sudo fc-cache -fv > /dev/null 2>&1
|
|
|
|
echo "Installed fonts:"
|
|
fc-list | grep -i fontin || echo "WARNING: Fonts not found after install!"
|
|
|
|
popd > /dev/null
|
|
rm -rf "$FONT_TMP"
|
|
}
|
|
|
|
# ── Step 3: Generate training images ─────────────────────────────────────────
|
|
|
|
generate_training_images() {
|
|
echo ""
|
|
echo "── Step 3: Generating training images ──"
|
|
|
|
mkdir -p "$WORK_DIR/ground-truth" "$WORK_DIR/output"
|
|
cd "$WORK_DIR"
|
|
|
|
# Copy eng.traineddata from project
|
|
if [ ! -f "$WORK_DIR/eng.traineddata" ]; then
|
|
cp "$TESSDATA_DIR/eng.traineddata" "$WORK_DIR/"
|
|
fi
|
|
|
|
# Extract LSTM model for fine-tuning
|
|
if [ ! -f "$WORK_DIR/eng.lstm" ]; then
|
|
echo "Extracting LSTM model from eng.traineddata..."
|
|
combine_tessdata -e eng.traineddata eng.lstm
|
|
fi
|
|
|
|
# Discover available Fontin font names
|
|
echo "Available Fontin fonts:"
|
|
text2image --list_available_fonts --fonts_dir "$FONT_DIR" 2>/dev/null | grep -i fontin || true
|
|
|
|
# Read font names into a file (avoids bash array portability issues)
|
|
# text2image output format is " 6: Fontin Bold" — extract font name after first colon
|
|
local FONT_LIST="$WORK_DIR/font_list.txt"
|
|
text2image --list_available_fonts --fonts_dir "$FONT_DIR" 2>/dev/null \
|
|
| grep -i fontin \
|
|
| cut -d: -f2- \
|
|
| sed 's/^[[:space:]]*//;s/[[:space:]]*$//' \
|
|
> "$FONT_LIST" 2>/dev/null || true
|
|
|
|
if [ ! -s "$FONT_LIST" ]; then
|
|
echo "WARNING: No Fontin fonts auto-detected. Trying common names..."
|
|
cat > "$FONT_LIST" <<'NAMES'
|
|
Fontin Bold
|
|
Fontin Medium
|
|
Fontin Medium Italic
|
|
Fontin SmallCaps, Medium
|
|
NAMES
|
|
fi
|
|
|
|
echo "Will generate images for:"
|
|
cat "$FONT_LIST"
|
|
|
|
local count=0
|
|
for exp in -1 0 1; do
|
|
while IFS= read -r font_name; do
|
|
[ -z "$font_name" ] && continue
|
|
local safe
|
|
safe=$(echo "$font_name" | tr ' ' '-' | tr '[:upper:]' '[:lower:]')
|
|
local outbase="ground-truth/poe2.${safe}.exp${exp}"
|
|
|
|
echo " Generating: $outbase (font='$font_name', exposure=$exp)"
|
|
if text2image \
|
|
--text "$TRAINING_TEXT" \
|
|
--outputbase "$outbase" \
|
|
--font "$font_name" \
|
|
--fonts_dir "$FONT_DIR" \
|
|
--ptsize 16 \
|
|
--xsize 3600 \
|
|
--ysize 480 \
|
|
--char_spacing 0.0 \
|
|
--exposure "$exp" \
|
|
--leading 32 2>&1; then
|
|
count=$((count + 1))
|
|
else
|
|
echo " WARNING: Failed for font '$font_name' exposure=$exp, skipping"
|
|
fi
|
|
done < "$FONT_LIST"
|
|
done
|
|
|
|
echo "Generated $count training image sets."
|
|
|
|
if [ "$count" -eq 0 ]; then
|
|
echo "ERROR: No training images generated. Check font installation."
|
|
return 1
|
|
fi
|
|
}
|
|
|
|
# ── Step 4: Generate LSTMF files ─────────────────────────────────────────────
|
|
|
|
generate_lstmf() {
|
|
echo ""
|
|
echo "── Step 4: Generating .lstmf training files ──"
|
|
|
|
cd "$WORK_DIR"
|
|
|
|
local count=0
|
|
for tif in ground-truth/*.tif; do
|
|
[ -f "$tif" ] || continue
|
|
local base="${tif%.tif}"
|
|
|
|
if [ -f "${base}.lstmf" ]; then
|
|
echo " Skipping (exists): ${base}.lstmf"
|
|
count=$((count + 1))
|
|
continue
|
|
fi
|
|
|
|
echo " Processing: $tif"
|
|
if tesseract "$tif" "$base" --psm 6 lstm.train 2>&1; then
|
|
count=$((count + 1))
|
|
else
|
|
echo " WARNING: Failed to process $tif"
|
|
fi
|
|
done
|
|
|
|
# Create training file list
|
|
: > training_files.txt
|
|
for f in ground-truth/*.lstmf; do
|
|
[ -f "$f" ] && echo "$f" >> training_files.txt
|
|
done
|
|
local total
|
|
total=$(wc -l < training_files.txt)
|
|
|
|
echo "Created $total .lstmf files."
|
|
|
|
if [ "$total" -eq 0 ]; then
|
|
echo "ERROR: No .lstmf files generated."
|
|
return 1
|
|
fi
|
|
}
|
|
|
|
# ── Step 5: Fine-tune LSTM ───────────────────────────────────────────────────
|
|
|
|
fine_tune() {
|
|
echo ""
|
|
echo "── Step 5: Fine-tuning LSTM (max $MAX_ITERATIONS iterations) ──"
|
|
|
|
cd "$WORK_DIR"
|
|
|
|
lstmtraining \
|
|
--continue_from eng.lstm \
|
|
--traineddata eng.traineddata \
|
|
--train_listfile training_files.txt \
|
|
--model_output output/poe2 \
|
|
--max_iterations "$MAX_ITERATIONS" \
|
|
--target_error_rate "$TARGET_ERROR" \
|
|
--debug_interval -1
|
|
|
|
echo "Fine-tuning complete."
|
|
}
|
|
|
|
# ── Step 6: Package model ────────────────────────────────────────────────────
|
|
|
|
package_model() {
|
|
echo ""
|
|
echo "── Step 6: Packaging poe2.traineddata ──"
|
|
|
|
cd "$WORK_DIR"
|
|
|
|
# Find the checkpoint file
|
|
local checkpoint=""
|
|
if [ -f "output/poe2_checkpoint" ]; then
|
|
checkpoint="output/poe2_checkpoint"
|
|
elif [ -f "output/poe2checkpoint" ]; then
|
|
checkpoint="output/poe2checkpoint"
|
|
else
|
|
# Find most recent checkpoint
|
|
checkpoint=$(ls -t output/poe2* 2>/dev/null | head -1)
|
|
fi
|
|
|
|
if [ -z "$checkpoint" ] || [ ! -f "$checkpoint" ]; then
|
|
echo "ERROR: No checkpoint found in output/"
|
|
ls -la output/ 2>/dev/null || true
|
|
return 1
|
|
fi
|
|
|
|
echo "Using checkpoint: $checkpoint"
|
|
|
|
lstmtraining --stop_training \
|
|
--continue_from "$checkpoint" \
|
|
--traineddata eng.traineddata \
|
|
--model_output output/poe2.traineddata
|
|
|
|
# Copy to project tessdata
|
|
cp output/poe2.traineddata "$TESSDATA_DIR/poe2.traineddata"
|
|
|
|
echo "Model saved to: $TESSDATA_DIR/poe2.traineddata"
|
|
ls -lh "$TESSDATA_DIR/poe2.traineddata"
|
|
}
|
|
|
|
# ── Step 7: Quick validation ─────────────────────────────────────────────────
|
|
|
|
validate() {
|
|
echo ""
|
|
echo "── Step 7: Quick validation ──"
|
|
|
|
cd "$WORK_DIR"
|
|
|
|
# Create a small test file
|
|
printf 'Quality: +20%%\nAdds 12 to 24 Fire Damage\n+45%% to Fire Resistance\n' > /tmp/poe2_test.txt
|
|
|
|
# Get first available fontin font name
|
|
local test_font
|
|
test_font=$(text2image --list_available_fonts --fonts_dir "$FONT_DIR" 2>/dev/null | grep -i fontin | head -1 | sed 's/^[[:space:]]*//;s/[[:space:]]*$//')
|
|
|
|
if [ -z "$test_font" ]; then
|
|
echo "Skipping validation (no Fontin font available)"
|
|
return 0
|
|
fi
|
|
|
|
if ! text2image \
|
|
--text /tmp/poe2_test.txt \
|
|
--outputbase /tmp/poe2_test \
|
|
--font "$test_font" \
|
|
--fonts_dir "$FONT_DIR" \
|
|
--ptsize 16 --xsize 1200 --ysize 200 \
|
|
--exposure 0 --leading 32 2>&1; then
|
|
echo "Skipping validation (could not generate test image)"
|
|
return 0
|
|
fi
|
|
|
|
echo "=== eng model ==="
|
|
tesseract /tmp/poe2_test.tif stdout -l eng --psm 6 2>/dev/null || true
|
|
|
|
echo ""
|
|
echo "=== poe2 model ==="
|
|
TESSDATA_PREFIX="$TESSDATA_DIR" tesseract /tmp/poe2_test.tif stdout -l poe2 --psm 6 2>/dev/null || {
|
|
tesseract /tmp/poe2_test.tif stdout --tessdata-dir "$TESSDATA_DIR" -l poe2 --psm 6 2>/dev/null || true
|
|
}
|
|
|
|
rm -f /tmp/poe2_test.txt /tmp/poe2_test.tif /tmp/poe2_test.box
|
|
|
|
echo ""
|
|
echo "Compare the outputs above — poe2 should be more accurate on Fontin text."
|
|
}
|
|
|
|
# ── Main ─────────────────────────────────────────────────────────────────────
|
|
|
|
main() {
|
|
install_tesseract_tools
|
|
install_fonts
|
|
generate_training_images
|
|
generate_lstmf
|
|
fine_tune
|
|
package_model
|
|
validate
|
|
|
|
echo ""
|
|
echo "=== Done! ==="
|
|
echo "poe2.traineddata has been copied to: $TESSDATA_DIR/"
|
|
echo ""
|
|
echo "Next steps:"
|
|
echo " 1. Build the daemon: dotnet build tools/OcrDaemon -c Release"
|
|
echo " 2. Start the bot: npx tsx src/index.ts"
|
|
echo " 3. Test OCR quality in the dashboard"
|
|
}
|
|
|
|
main "$@"
|