Update download_resource.bat

download_resource.bat

@@ -1,22 +1,72 @@
 @echo off
 setlocal enabledelayedexpansion
 
+REM 错误处理
+:ERROR_HANDLER
+IF %ERRORLEVEL% NEQ 0 (
+    echo.
+    echo ************************************************************
+    echo *                    下载任务被中断                         *
+    echo *              请检查网络连接后重新运行脚本                 *
+    echo *          如果问题持续存在，请查看以上错误信息             *
+    echo ************************************************************
+    exit /b 1
+)
+
+REM 设置当前目录
+SET PROJECT_DIR=%~dp0
+cd %PROJECT_DIR%
+
+REM 捕获Ctrl+C中断
+REM 在Windows中，可以通过注册一个处理函数来捕获Ctrl+C
+IF NOT DEFINED CAUGHT_CTRL_C (
+    SET CAUGHT_CTRL_C=1
+    CALL :REGISTER_CTRL_C_HANDLER
+)
+
+REM 定义文件路径
+SET "NEURAL_HAIRCUT_FOLDER_PATH=resource\NeuralHaircut"
+SET "PIXIE_PATH=resource\NeuralHaircut\PIXIE\pixie_data\pixie_data"
+SET "MATTE_PATH=resource\Matte-Anything\model.pth"
+SET "SAM_PATH=resource\Matte-Anything\sam_vit_h_4b8939.pth"
+SET "DINO_PATH=resource\Matte-Anything\groundingdino_swint_ogc.pth"
+SET "OPENPOSE_PATH=resource\openpose\models\pose\coco\pose_iter_584000.caffemodel"
+SET "HYPERIQA_PATH=resource\hyperIQA\pretrained\hyperIQA.pth"
+
 REM 清理旧的资源目录
 IF EXIST "resource" (
     echo 清理旧的资源目录...
     rmdir /s /q "resource"
 )
-else (
+
-    echo 资源目录不存在，创建资源目录...
+REM 创建resource目录结构
-    mkdir resource
+echo 正在创建resource目录结构...
+mkdir resource 2>nul
+mkdir resource\NeuralHaircut 2>nul
+mkdir resource\NeuralHaircut\PIXIE 2>nul
+mkdir resource\Matte-Anything 2>nul
+mkdir resource\Matte-Anything\pretrained 2>nul
+mkdir resource\openpose 2>nul
+mkdir resource\hyperIQA 2>nul
+mkdir resource\hyperIQA\pretrained 2>nul
+
+REM 验证已有文件
+echo 正在检查已有资源文件...
+SET MISSING_FILES=0
+SET /A TOTAL_FILES=7
+SET /A EXISTING_FILES=0
+SET MISSING_LIST=
+
+REM Neural Haircut Diffusion Prior 检查
+IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
+    echo [√] Neural Haircut Diffusion Prior已存在
+    SET /A EXISTING_FILES+=1
+) ELSE (
+    echo [×] 缺少Neural Haircut Diffusion Prior
+    SET /A MISSING_FILES+=1
+    SET MISSING_LIST=!MISSING_LIST!Neural Haircut Diffusion Prior, 
 )
 
-
-
-REM 设置当前目录
-SET PROJECT_DIR=%~dp0
-cd %PROJECT_DIR%
-
 REM 检查7-Zip是否安装
 where 7z >nul 2>nul
 IF %ERRORLEVEL% NEQ 0 (
@@ -89,22 +139,26 @@ REM Download Neural Haircut files
 echo 正在下载Neural Haircut文件...
 cd %PROJECT_DIR%/resource/NeuralHaircut
 python -m gdown --folder https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8
+CALL :ERROR_HANDLER
 
 REM 更新 diffusion_prior 模型
 cd %PROJECT_DIR%/resource/NeuralHaircut/pretrained_models/diffusion_prior
 echo 正在下载Neural Haircut扩散先验模型文件...
 python -m gdown 1_9EOUXHayKiGH5nkrayncln3d6m1uV7f
+CALL :ERROR_HANDLER
 
 REM 下载 PIXIE 文件
 cd %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
 echo 正在下载Neural HaircutPIXIE模型文件...
 python -m gdown 1mPcGu62YPc4MdkT8FFiOCP629xsENHZf && tar -xvzf pixie_data.tar.gz ./ && rm pixie_data.tar.gz
+CALL :ERROR_HANDLER
 
 REM 下载 hyperIQA 模型
 cd %PROJECT_DIR%/resource/hyperIQA/pretrained
 echo 正在下载hyperIQA模型文件...
 python -m gdown 1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE
 cd %PROJECT_DIR%
+CALL :ERROR_HANDLER
 
 REM Download Matte-Anything files
 cd %PROJECT_DIR%/resource/Matte-Anything/pretrained
@@ -114,15 +168,63 @@ echo 正在下载Matte-AnythingGroundingDINO模型文件...
 wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
 echo 正在下载Matte-Anything模型文件...
 python -m gdown 1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW
+CALL :ERROR_HANDLER
 
 REM Download OpenPose files
 cd %PROJECT_DIR%/resource/openpose
 echo 正在下载OpenPose模型文件...
 python -m gdown 1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV && tar -xvzf models.tar.gz && rm models.tar.gz
+CALL :ERROR_HANDLER
 
 cd %PROJECT_DIR%
 echo 所有模型文件下载完成！
 
+echo.
+echo 已存在 %EXISTING_FILES%/%TOTAL_FILES% 个模型文件
+IF %MISSING_FILES% GTR 0 (
+    echo 需要下载 %MISSING_FILES% 个模型文件
+    echo 缺少的模型: !MISSING_LIST:~0,-2!
+    echo.
+    echo ************************************************************
+    echo *                  开始下载所需资源文件                     *
+    echo *            下载时间取决于您的网络连接速度                 *
+    echo *                    请耐心等待...                         *
+    echo ************************************************************
+) ELSE (
+    echo 所有模型文件已存在，无需下载
+    exit /b 0
+)
+
+IF %VERIFY_FAILED% GTR 0 (
+    echo.
+    echo 文件完整性验证失败：%VERIFY_FAILED% 个文件缺失或不完整
+    echo 请重新运行脚本下载缺失文件
+    exit /b 1
+) ELSE (
+    echo.
+    echo 文件完整性验证成功！
+)
+
+GOTO :EOF
+
+:REGISTER_CTRL_C_HANDLER
+echo 注册Ctrl+C处理函数...
+FOR /F "tokens=2 delims=[]" %%a IN ('ver') DO FOR /F "tokens=2,3 delims=. " %%b IN ("%%a") DO SET /A WINVER=%%b%%c
+IF %WINVER% GEQ 100 (
+    REM Windows 10及以上版本
+    powershell -Command "$handler = [Console]::CancelKeyPress.ToString(); if ($handler -notmatch 'Ctrl\+C') { [Console]::TreatControlCAsInput = $true }"
+)
+
+:CTRL_C_HANDLER
+echo.
+echo ************************************************************
+echo *                    下载任务被用户中断                     *
+echo *                  请重新运行脚本继续下载                   *
+echo ************************************************************
+exit /b 1
+
+:EOF
+