hongchen/Floorplan/run_pipeline.sh

#!/bin/bash

# ============================================================================
# Floorplan Pipeline 一键执行脚本
# ============================================================================
# 功能：依次执行 extract_initial_mask -> inference_refine_mask -> pipeline -> vis
# 环境：conda activate floorplan_env
# 用法：./run_pipeline.sh <图片名称>
# 示例：./run_pipeline.sh SG-n6nV8B2oW95
# ============================================================================

# 记录开始时间
START_TIME=$(date +%s)

# ============================================================================
# 内存和显存监控变量
# ============================================================================
MONITOR_DIR="/tmp/floorplan_monitor_$$"
mkdir -p "$MONITOR_DIR"
echo "0" > "$MONITOR_DIR/max_ram"
echo "0" > "$MONITOR_DIR/max_vram"
MONITOR_PID=""

# 获取所有 python 进程的总 RAM 占用 (MB)
get_python_ram_usage() {
    local total=0
    for pid in $(pgrep -f "python.*\.py" 2>/dev/null); do
        if [ -f "/proc/$pid/status" ]; then
            local vmrss=$(grep VmRSS /proc/$pid/status 2>/dev/null | awk '{print $2}')
            if [ -n "$vmrss" ]; then
                total=$((total + vmrss / 1024))
            fi
        fi
    done
    echo "$total"
}

# 获取 GPU 显存占用 (MB)
get_vram_usage() {
    local used=$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | awk '{sum+=$1} END {print sum+0}')
    echo "$used"
}

# 后台监控进程（使用子 shell）
monitor_memory() {
    local max_ram_file="$MONITOR_DIR/max_ram"
    local max_vram_file="$MONITOR_DIR/max_vram"
    local running_flag="$MONITOR_DIR/running"
    echo "1" > "$running_flag"
    while [ "$(cat "$running_flag" 2>/dev/null)" = "1" ]; do
        local ram=$(get_python_ram_usage)
        local vram=$(get_vram_usage)
        # 更新最大值（空值默认为 0）
        local cur_max_ram=$(cat "$max_ram_file" 2>/dev/null || echo "0")
        local cur_max_vram=$(cat "$max_vram_file" 2>/dev/null || echo "0")
        cur_max_ram=${cur_max_ram:-0}
        cur_max_vram=${cur_max_vram:-0}
        [ -n "$ram" ] && [ "$ram" -gt "$cur_max_ram" ] 2>/dev/null && echo "$ram" > "$max_ram_file"
        [ -n "$vram" ] && [ "$vram" -gt "$cur_max_vram" ] 2>/dev/null && echo "$vram" > "$max_vram_file"
        sleep 1
    done
}

echo "=============================================="
echo "  Floorplan Pipeline 一键处理"
echo "=============================================="
echo ""

# 检查参数
if [ -z "$1" ]; then
    echo "用法：./run_pipeline.sh <图片名称>"
    echo "示例：./run_pipeline.sh SG-n6nV8B2oW95"
    exit 1
fi

IMAGE_NAME="$1"

# 1. 激活 conda 环境
echo "🔧 激活 conda 环境：floorplan_env"
source /home/dage/Project/Miniconda-25-1/etc/profile.d/conda.sh
conda activate floorplan_env

if [ $? -ne 0 ]; then
    echo "❌ 激活 conda 环境失败，请检查环境名称是否正确"
    exit 1
fi

echo "✅ 环境激活成功"
echo ""

# 2. 检查图片是否存在
if [ ! -f "temp_data/${IMAGE_NAME}" ] && [ ! -f "temp_data/${IMAGE_NAME}.png" ] && [ ! -f "temp_data/${IMAGE_NAME}.jpg" ]; then
    echo "❌ 图片不存在：temp_data/${IMAGE_NAME}*"
    echo "   请确认图片名称（可带扩展名或不带）"
    exit 1
fi

echo "📁 处理图片：${IMAGE_NAME}"
echo ""

# 启动后台监控
echo "📊 内存/显存监控已启动"
echo ""
monitor_memory &
MONITOR_PID=$!

# ============================================================================
# Step 1: 提取 initial mask
# ============================================================================
echo "=============================================="
echo "  Step 1/4: 提取 Initial Mask"
echo "=============================================="
python extract_initial_mask.py "${IMAGE_NAME}"
if [ $? -ne 0 ]; then
    echo "❌ Step 1 失败，终止流程"
    exit 1
fi
echo ""

# 清理显存
echo "🧹 清理显存..."
python -c "import torch; torch.cuda.empty_cache()"
echo ""

# ============================================================================
# Step 2: 推理 refine mask
# ============================================================================
echo "=============================================="
echo "  Step 2/4: 推理 Refine Mask"
echo "=============================================="
BASE_NAME="${IMAGE_NAME%.*}"  # 去掉扩展名
python inference_refine_mask.py "${BASE_NAME}"
if [ $? -ne 0 ]; then
    echo "❌ Step 2 失败，终止流程"
    exit 1
fi
echo ""

# 清理显存
echo "🧹 清理显存..."
python -c "import torch; torch.cuda.empty_cache()"
echo ""

# ============================================================================
# Step 3: Pipeline 处理生成 JSON
# ============================================================================
echo "=============================================="
echo "  Step 3/4: Pipeline 处理生成 JSON"
echo "=============================================="
python pipeline.py "${BASE_NAME}"
if [ $? -ne 0 ]; then
    echo "❌ Step 3 失败，终止流程"
    exit 1
fi
echo ""

# ============================================================================
# Step 4: 可视化
# ============================================================================
echo "=============================================="
echo "  Step 4/4: 可视化"
echo "=============================================="
python vis.py "${BASE_NAME}"
if [ $? -ne 0 ]; then
    echo "❌ Step 4 失败，终止流程"
    exit 1
fi
echo ""

# ============================================================================
# 完成
# ============================================================================

# 停止监控进程
echo "0" > "$MONITOR_DIR/running"
sleep 1
kill $MONITOR_PID 2>/dev/null
wait $MONITOR_PID 2>/dev/null

# 读取最大值
MAX_RAM_MB=$(cat "$MONITOR_DIR/max_ram" 2>/dev/null)
MAX_VRAM_MB=$(cat "$MONITOR_DIR/max_vram" 2>/dev/null)
MAX_RAM_MB=${MAX_RAM_MB:-0}
MAX_VRAM_MB=${MAX_VRAM_MB:-0}

# 清理监控目录
rm -rf "$MONITOR_DIR"

echo "=============================================="
echo "  ✅ 全部处理完成！"
echo "=============================================="
echo ""
echo "📁 结果位置：temp_data/${BASE_NAME}/"
ls -1 "temp_data/${BASE_NAME}/"

# 计算并显示运行时间
END_TIME=$(date +%s)
ELAPSED=$((END_TIME - START_TIME))

# 格式化时间显示（超过 60 秒显示为分钟）
if [ $ELAPSED -ge 60 ]; then
    MINUTES=$((ELAPSED / 60))
    SECONDS=$((ELAPSED % 60))
    echo "⏱️  总耗时：${MINUTES}分${SECONDS}秒"
else
    echo "⏱️  总耗时：${ELAPSED}秒"
fi

# 显示内存和显存使用情况
echo ""
echo "=============================================="
echo "  📊 资源使用情况"
echo "=============================================="
echo "  最大内存占用 (RAM): ${MAX_RAM_MB} MB"
echo "  最大显存占用 (VRAM): ${MAX_VRAM_MB} MB"
echo "=============================================="
echo ""