#!/bin/bash
# Quick training script for Fish Action Classification

set -e
cd /home/ubuntu/projects/FishAction/slowfast

# # Configuration
CONFIG_FILE="../configs/fish_action_SLOWFAST_8x8_R50.yaml"
PRETRAINED_MODEL="checkpoints/SLOWFAST_8x8_R50.pkl"
# PRETRAINED_URL="https://dl.fbaipublicfiles.com/pyslowfast/model_zoo/kinetics400/SLOWFAST_8x8_R50.pkl"

# Colors for output
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

echo -e "${GREEN}Fish Action Classification Training${NC}"
echo "=========================================="

# Check if pretrained model exists
# if [ ! -f "$PRETRAINED_MODEL" ]; then
#     echo -e "${YELLOW}Pretrained model not found. Downloading...${NC}"
#     mkdir -p checkpoints
#     wget "$PRETRAINED_URL" -O "$PRETRAINED_MODEL"
#     echo -e "${GREEN}Download complete!${NC}"
# else
#     echo -e "${GREEN}Pretrained model found: $PRETRAINED_MODEL${NC}"
# fi

# Check if config exists
if [ ! -f "$CONFIG_FILE" ]; then
    echo -e "${YELLOW}Error: Config file not found: $CONFIG_FILE${NC}"
    exit 1
fi

# Check if CSV files exist
CSV_DIR="/home/ubuntu/data/fish/fish_action_videos"


echo ""
echo "Starting training..."
echo "Config: $CONFIG_FILE"
echo "Pretrained: $PRETRAINED_MODEL"
echo ""

# Run training

python tools/run_net.py \
    --cfg "$CONFIG_FILE" \
    TRAIN.CHECKPOINT_FILE_PATH "$PRETRAINED_MODEL" \
    TRAIN.CHECKPOINT_TYPE caffe2 \
    NUM_GPUS ${NUM_GPUS:-1}

echo ""
echo -e "${GREEN}Training completed!${NC}"
echo "Checkpoints saved to: checkpoints/fish_action/"