Update

2025-02-16 22:40:27 +08:00 · 2025-02-16 22:40:27 +08:00 · df2080f58b
commit df2080f58b
parent feb2269e64
3 changed files with 106 additions and 48 deletions
--- a/download_resource.bat
+++ b/download_resource.bat
@ -4,6 +4,14 @@ setlocal enabledelayedexpansion
 SET PROJECT_DIR=%~dp0
 cd %PROJECT_DIR%

+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"

 echo ======================================= 基础环境配置 ==============================================
 echo.
@ -79,45 +87,102 @@ echo.
 echo [1/8] 正在下载Neural Haircut文件...
 mkdir %PROJECT_DIR%/resource/NeuralHaircut
 cd %PROJECT_DIR%/resource/NeuralHaircut
-python -m gdown --folder "https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8"
-cd %PROJECT_DIR%
+IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
+    IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\strand_ckpt.pt" (
+        echo Neural Haircut文件已存在，跳过下载...
+    ) ELSE (
+        echo Neural Haircut Strand Prior模型不完整，重新下载...
+        python -m gdown --folder "https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8"
+    )
+) ELSE (
+    echo Neural Haircut Diffusion Prior模型不完整，重新下载...
+    python -m gdown --folder "https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8"
+)

-echo [2/8] 正在下载Neural Haircut扩散先验模型文件...
-cd %PROJECT_DIR%/resource/NeuralHaircut/pretrained_models/diffusion_prior
-python -m gdown "1_9EOUXHayKiGH5nkrayncln3d6m1uV7f"
-cd %PROJECT_DIR%
+REM 检查下载后的文件完整性
+IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
+    echo 错误：Neural Haircut Diffusion Prior模型下载失败
+    exit /b 1
+)
+IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\strand_ckpt.pt" (
+    echo 错误：Neural Haircut Strand Prior模型下载失败
+    exit /b 1
+)

-echo [3/8] 正在下载Neural HaircutPIXIE模型文件...
-mkdir %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
-cd %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
-python -m gdown "1mPcGu62YPc4MdkT8FFiOCP629xsENHZf" && tar -xvzf pixie_data.tar.gz ./ && rm pixie_data.tar.gz
-cd %PROJECT_DIR%
+echo [2/8] 正在下载Neural HaircutPIXIE模型文件...
+IF EXIST "%PIXIE_PATH%" (
+    echo PIXIE模型已存在，跳过下载...
+) ELSE (
+    mkdir %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
+    cd %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
+    python -m gdown "1mPcGu62YPc4MdkT8FFiOCP629xsENHZf" && tar -xvzf pixie_data.tar.gz ./ && rm pixie_data.tar.gz
+    cd %PROJECT_DIR%
+)

-echo [4/8] 正在下载hyperIQA模型文件...
+echo [3/8] 正在下载hyperIQA模型文件...
 mkdir %PROJECT_DIR%/resource/hyperIQA/pretrained
 cd %PROJECT_DIR%/resource/hyperIQA/pretrained
-python -m gdown "1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE"
+IF EXIST "hyperIQA.pth" (
+    echo hyperIQA模型已存在，跳过下载...
+) ELSE (
+    python -m gdown "1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE"
+    IF NOT EXIST "hyperIQA.pth" (
+        echo 错误：hyperIQA模型下载失败
+        exit /b 1
+    )
+)
 cd %PROJECT_DIR%

-echo [5/8] 正在下载Matte-Anything扩散先验模型文件...
+echo [4/8] 正在下载Matte-Anything扩散先验模型文件...
 mkdir %PROJECT_DIR%/resource/Matte-Anything/pretrained
 cd %PROJECT_DIR%/resource/Matte-Anything/pretrained
-wget "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
-echo [6/8] 正在下载Matte-AnythingGroundingDINO模型文件...
-cd %PROJECT_DIR%/resource/Matte-Anything/pretrained
-wget "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
-echo [7/8] 正在下载Matte-Anything模型文件...
-python -m gdown "1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW"
+IF EXIST "sam_vit_h_4b8939.pth" (
+    echo SAM模型已存在，跳过下载...
+) ELSE (
+    wget "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+    IF NOT EXIST "sam_vit_h_4b8939.pth" (
+        echo 错误：SAM模型下载失败
+        exit /b 1
+    )
+)
+
+echo [5/8] 正在下载Matte-AnythingGroundingDINO模型文件...
+IF EXIST "groundingdino_swint_ogc.pth" (
+    echo GroundingDINO模型已存在，跳过下载...
+) ELSE (
+    wget "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
+    IF NOT EXIST "groundingdino_swint_ogc.pth" (
+        echo 错误：GroundingDINO模型下载失败
+        exit /b 1
+    )
+)
+
+echo [6/8] 正在下载Matte-Anything模型文件...
+IF EXIST "model.pth" (
+    echo Matte-Anything模型已存在，跳过下载...
+) ELSE (
+    python -m gdown "1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW"
+    IF NOT EXIST "model.pth" (
+        echo 错误：Matte-Anything模型下载失败
+        exit /b 1
+    )
+)
 cd %PROJECT_DIR%

-echo [8/8] 正在下载OpenPose模型文件...
+echo [7/8] 正在下载OpenPose模型文件...
 mkdir %PROJECT_DIR%/resource/openpose/models/pose/coco
 cd %PROJECT_DIR%/resource/openpose/models/pose/coco
-python -m gdown "1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV" && tar -xvzf models.tar.gz && rm models.tar.gz
+IF EXIST "pose_iter_584000.caffemodel" (
+    echo OpenPose模型已存在，跳过下载...
+) ELSE (
+    python -m gdown "1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV" && tar -xvzf models.tar.gz && rm models.tar.gz
+    IF NOT EXIST "pose_iter_584000.caffemodel" (
+        echo 错误：OpenPose模型下载失败
+        exit /b 1
+    )
+)
 cd %PROJECT_DIR%

-echo 所有模型文件下载完成！
-



--- a/install.bat
+++ b/install.bat
@ -159,20 +159,31 @@ cd %PROJECT_DIR%

 REM 检查resource目录
 IF NOT EXIST "%PROJECT_DIR%\resource" (
-    echo ERROR: resource directory not found
-    echo Please ensure the resource directory exists with required files:
+    echo 错误：未找到resource目录
+    echo 请确保resource目录存在且包含以下必需文件：
    echo.
    echo resource/
    echo ├── NeuralHaircut/
-    echo │   ├── diffusion_prior/
-    echo │   │   └── dif_ckpt.pt
-    echo │   ├── strand_prior/
+    echo │   ├── pretrained_models/
+    echo │   │   ├── diffusion_prior/
+    echo │   │   │   └── dif_ckpt.pt
+    echo │   │   └── strand_prior/
    echo │   │       └── strand_ckpt.pt
    echo │   └── PIXIE/
    echo │       └── pixie_data/
+    echo │           └── pixie_data/
    echo ├── Matte-Anything/
+    echo │   ├── model.pth
+    echo │   ├── sam_vit_h_4b8939.pth
+    echo │   └── groundingdino_swint_ogc.pth
    echo ├── openpose/
+    echo │   └── models/
+    echo │       └── pose/
+    echo │           └── coco/
+    echo │               └── pose_iter_584000.caffemodel
    echo └── hyperIQA/
+    echo     └── pretrained/
+    echo         └── hyperIQA.pth
    exit /b 1
 )

--- a/run.bat
+++ b/run.bat
@ -340,21 +340,3 @@ IF %ERRORLEVEL% NEQ 0 (
    echo ERROR: Failed to run command
    exit /b 1
 )
-
-REM 检查模型文件
-IF NOT EXIST "%PROJECT_DIR%\ext\NeuralHaircut\diffusion_prior\dif_ckpt.pt" (
-    echo 错误：未找到Neural Haircut模型，请先运行install.bat
-    exit /b 1
-)
-
-REM 检查输入数据
-IF NOT EXIST "%DATA_PATH%" (
-    echo 错误：未找到数据目录 %DATA_PATH%
-    exit /b 1
-)
-
-REM 检查CUDA设备
-nvidia-smi >nul 2>nul || (
-    echo 错误：未找到NVIDIA GPU或驱动未正确安装
-    exit /b 1
-)