diff --git a/download_resource.bat b/download_resource.bat index 7fbb18b..2e7fadf 100644 --- 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 所有模型文件下载完成! - diff --git a/install.bat b/install.bat index cd0e7a9..21d1fa2 100644 --- 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 │ │ └── strand_ckpt.pt + 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 ) diff --git a/run.bat b/run.bat index 31a8799..81aa267 100644 --- 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 -)