hongchen/Floorplan/run_pipeline_simple.sh

#!/bin/bash

# ============================================================================
# Floorplan Pipeline Docker 一键执行脚本
# ============================================================================
# 功能：自动启动容器并依次执行 extract_initial_mask -> inference_refine_mask ->
#       pipeline -> vis
# 镜像：floorplan_env_full:latest (自带 CUDA 库)
# 用法：./run_pipeline_docker.sh <图片名称>
# 示例：./run_pipeline_docker.sh SG-n6nV8B2oW95
# ============================================================================

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

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

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

IMAGE_NAME="$1"

# 获取当前工作目录的绝对路径
WORK_DIR=$(pwd)

# Docker 镜像和容器名称
IMAGE_NAME_DOCKER="floorplan_env_full:latest"
CONTAINER_NAME="floorplan_container"

# 获取实际图片文件名（支持多种格式）
get_image_file() {
    local img_name="$1"
    local img_dir="${WORK_DIR}/temp_data"

    # 检查是否是目录（已处理的图片）
    if [ -d "${img_dir}/${img_name}" ]; then
        # 在目录中查找原图
        for ext in png jpg jpeg; do
            if [ -f "${img_dir}/${img_name}/${img_name}.${ext}" ]; then
                echo "${img_name}/${img_name}.${ext}"
                return 0
            fi
        done
    fi

    # 检查是否是文件
    if [ -f "${img_dir}/${img_name}" ]; then
        echo "${img_name}"
        return 0
    elif [ -f "${img_dir}/${img_name}.png" ]; then
        echo "${img_name}.png"
        return 0
    elif [ -f "${img_dir}/${img_name}.jpg" ]; then
        echo "${img_name}.jpg"
        return 0
    fi

    return 1
}

ACTUAL_IMAGE_FILE=$(get_image_file "${IMAGE_NAME}")
if [ -z "${ACTUAL_IMAGE_FILE}" ]; then
    echo "❌ 图片不存在：${WORK_DIR}/temp_data/${IMAGE_NAME}*"
    echo "   请确认图片名称（可带扩展名或不带）"
    exit 1
fi

# 检查并启动容器
if docker ps --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
    echo "✅ 容器 '${CONTAINER_NAME}' 已在运行"
else
    echo "🚀 启动 Docker 容器 '${CONTAINER_NAME}'..."

    # 检查容器是否存在（已停止的）
    if docker ps -a --format '{{.Names}}' | grep -q "^${CONTAINER_NAME}$"; then
        # 容器存在，直接启动
        docker start ${CONTAINER_NAME}
    else
        # 容器不存在，创建新容器
        docker run -td \
            --gpus all \
            --name ${CONTAINER_NAME} \
            --ipc=host \
            -v "${WORK_DIR}":"${WORK_DIR}" \
            -w "${WORK_DIR}" \
            ${IMAGE_NAME_DOCKER}

        if [ $? -ne 0 ]; then
            echo "❌ 启动容器失败"
            exit 1
        fi
        echo "✅ 容器启动成功（镜像已自带 CUDA 库）"
    fi
fi
echo ""

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

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

# 清理显存
echo "🧹 清理显存..."
docker exec -w "${WORK_DIR}" ${CONTAINER_NAME} 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%.*}"  # 去掉扩展名
docker exec -w "${WORK_DIR}" ${CONTAINER_NAME} python inference_refine_mask.py "${BASE_NAME}"
if [ $? -ne 0 ]; then
    echo "❌ Step 2 失败，终止流程"
    exit 1
fi
echo ""

# 清理显存
echo "🧹 清理显存..."
docker exec -w "${WORK_DIR}" ${CONTAINER_NAME} python -c "import torch; torch.cuda.empty_cache()"
echo ""

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

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

# ============================================================================
# 完成
# ============================================================================
echo "=============================================="
echo "  ✅ 全部处理完成！"
echo "=============================================="
echo ""

# 修复文件权限（将容器内 root 创建的文件所有权改为宿主机用户）
echo "🔧 修复文件权限..."
docker exec ${CONTAINER_NAME} chown -R $(id -u):$(id -g) "${WORK_DIR}/temp_data/${BASE_NAME}"
echo ""

echo "📁 结果位置：temp_data/${BASE_NAME}/"
ls -1 "${WORK_DIR}/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