fix model-download-bugs

bcmi · Jan 25, 2024 · b6bd5c7 · b6bd5c7
1 parent 45abbfe
commit b6bd5c7
Show file tree

Hide file tree

Showing 53 changed files with 48 additions and 48 deletions.
diff --git a/.gitignore b/.gitignore
@@ -1 +1,14 @@
-./pyproject.toml
+./pyproject.toml
+
+/build
+/dist
+/*egg-info
+*README_cn.md
+zip_package.sh
+**/__pycache__
+**/pretrained_models
+make_package.sh
+**/shared_pretrained_models
+*template*
+**/*results
+**/developer_readme.md
diff --git a/libcom/utils/model_download.py b/libcom/utils/model_download.py
@@ -37,7 +37,7 @@ def download_entire_folder(folder_path):
 
 def download_file_from_network(file_name, save_dir):
     if not os.path.exists(save_dir):
-         os.mkdir(save_dir)
+        os.mkdir(save_dir)
     print('Try to download {} to {}'.format(file_name, save_dir))
     try:
         from huggingface_hub import hf_hub_download
@@ -51,8 +51,10 @@ def download_file_from_network(file_name, save_dir):
                                         cache_dir=save_dir, 
                                         revision='master')
     assert os.path.exists(file_path), 'Download {} failed, please try again'.format(file)
-    shutil.move(file_path, os.path.join(save_dir, file_name))
-
+    save_path = os.path.abspath(os.path.join(save_dir, file_name))
+    shutil.copyfile(os.path.abspath(file_path), save_path, follow_symlinks=True)
+    assert os.path.exists(save_path), 'Move file to {} failed, please try again'.format(save_path)
+    os.remove(os.path.realpath(file_path)) # delete the cache
 
 if __name__ == '__main__':
     file_list   = ['BargainNet.pth', 'SOPA.pth']
@@ -63,9 +65,4 @@ def download_file_from_network(file_name, save_dir):
 
     for folder in folder_list:
         folder_path = './pretrained_models/' + folder
-        download_entire_folder(folder_path)
-
-
-
-
-
+        download_entire_folder(folder_path)
diff --git a/setup.py b/setup.py
@@ -92,6 +92,7 @@ def gen_packages_items():
     return packages
 
 def get_ext_modules():
+    # if encounter compilation issues, please refer to https://github.com/HuiZeng/Image-Adaptive-3DLUT?tab=readme-ov-file#build.
     if torch.cuda.is_available():
         return CUDAExtension('trilinear', 
                              ['libcom/image_harmonization/source/trilinear_cpp/src/trilinear_cuda.cpp', 
@@ -107,7 +108,7 @@ def get_ext_modules():
     cur_dir = os.path.dirname(os.path.abspath(__file__))
     setup(
         name='libcom',
-        version='0.0.1.post8',
+        version='0.0.1.post9',
         description='Image Composition Toolbox',
         long_description=readme(),
         long_description_content_type='text/markdown',

diff --git a/tests/source/background/1.jpg b/tests/source/background/1.jpg
diff --git a/tests/source/background/17.jpg b/tests/source/background/17.jpg
diff --git a/tests/source/background/1bb218825e370c51_m02p5f1q_4a9db762_30.png b/tests/source/background/1bb218825e370c51_m02p5f1q_4a9db762_30.png
diff --git a/tests/source/background/35149c4895418040_m01yrx_e5b3a7bd_09.png b/tests/source/background/35149c4895418040_m01yrx_e5b3a7bd_09.png
diff --git a/tests/source/background/5b9e0751ca458cad_m0gjkl_40736b77_04.png b/tests/source/background/5b9e0751ca458cad_m0gjkl_40736b77_04.png
diff --git a/tests/source/background/8.jpg b/tests/source/background/8.jpg
diff --git a/tests/source/background/8351494f79f6bf8f_m03k3r_73176af9_08.png b/tests/source/background/8351494f79f6bf8f_m03k3r_73176af9_08.png
diff --git a/tests/source/background/f80eda2459853824_m09g1w_b2413ec8_11.png b/tests/source/background/f80eda2459853824_m09g1w_b2413ec8_11.png
diff --git a/tests/source/bbox/1.txt b/tests/source/bbox/1.txt
@@ -0,0 +1 @@
+1000 895 1480 1355
diff --git a/tests/source/bbox/17.txt b/tests/source/bbox/17.txt
@@ -0,0 +1 @@
+475 697 1275 1401
diff --git a/tests/source/bbox/1bb218825e370c51_m02p5f1q_4a9db762_30.txt b/tests/source/bbox/1bb218825e370c51_m02p5f1q_4a9db762_30.txt
diff --git a/tests/source/bbox/35149c4895418040_m01yrx_e5b3a7bd_09.txt b/tests/source/bbox/35149c4895418040_m01yrx_e5b3a7bd_09.txt
diff --git a/tests/source/bbox/5b9e0751ca458cad_m0gjkl_40736b77_04.txt b/tests/source/bbox/5b9e0751ca458cad_m0gjkl_40736b77_04.txt
diff --git a/tests/source/bbox/8.txt b/tests/source/bbox/8.txt
@@ -0,0 +1 @@
+1170 944 2331 3069
diff --git a/tests/source/bbox/8351494f79f6bf8f_m03k3r_73176af9_08.txt b/tests/source/bbox/8351494f79f6bf8f_m03k3r_73176af9_08.txt
diff --git a/tests/source/bbox/f80eda2459853824_m09g1w_b2413ec8_11.txt b/tests/source/bbox/f80eda2459853824_m09g1w_b2413ec8_11.txt
diff --git a/tests/source/composite/1.jpg b/tests/source/composite/1.jpg
diff --git a/tests/source/composite/17.jpg b/tests/source/composite/17.jpg
diff --git a/tests/source/composite/1bb218825e370c51_m02p5f1q_4a9db762_30.png b/tests/source/composite/1bb218825e370c51_m02p5f1q_4a9db762_30.png
diff --git a/tests/source/composite/35149c4895418040_m01yrx_e5b3a7bd_09.png b/tests/source/composite/35149c4895418040_m01yrx_e5b3a7bd_09.png
diff --git a/tests/source/composite/5b9e0751ca458cad_m0gjkl_40736b77_04.png b/tests/source/composite/5b9e0751ca458cad_m0gjkl_40736b77_04.png
diff --git a/tests/source/composite/8.jpg b/tests/source/composite/8.jpg
diff --git a/tests/source/composite/8351494f79f6bf8f_m03k3r_73176af9_08.png b/tests/source/composite/8351494f79f6bf8f_m03k3r_73176af9_08.png
diff --git a/tests/source/composite/f80eda2459853824_m09g1w_b2413ec8_11.png b/tests/source/composite/f80eda2459853824_m09g1w_b2413ec8_11.png
diff --git a/tests/source/composite_mask/1.png b/tests/source/composite_mask/1.png
diff --git a/tests/source/composite_mask/17.png b/tests/source/composite_mask/17.png
diff --git a/tests/source/composite_mask/1bb218825e370c51_m02p5f1q_4a9db762_30.png b/tests/source/composite_mask/1bb218825e370c51_m02p5f1q_4a9db762_30.png
diff --git a/tests/source/composite_mask/35149c4895418040_m01yrx_e5b3a7bd_09.png b/tests/source/composite_mask/35149c4895418040_m01yrx_e5b3a7bd_09.png
diff --git a/tests/source/composite_mask/5b9e0751ca458cad_m0gjkl_40736b77_04.png b/tests/source/composite_mask/5b9e0751ca458cad_m0gjkl_40736b77_04.png
diff --git a/tests/source/composite_mask/8.png b/tests/source/composite_mask/8.png
diff --git a/tests/source/composite_mask/8351494f79f6bf8f_m03k3r_73176af9_08.png b/tests/source/composite_mask/8351494f79f6bf8f_m03k3r_73176af9_08.png
diff --git a/tests/source/composite_mask/f80eda2459853824_m09g1w_b2413ec8_11.png b/tests/source/composite_mask/f80eda2459853824_m09g1w_b2413ec8_11.png
diff --git a/tests/source/foreground/1.jpg b/tests/source/foreground/1.jpg
diff --git a/tests/source/foreground/17.jpg b/tests/source/foreground/17.jpg
diff --git a/tests/source/foreground/1bb218825e370c51_m02p5f1q_4a9db762_30.png b/tests/source/foreground/1bb218825e370c51_m02p5f1q_4a9db762_30.png
diff --git a/tests/source/foreground/35149c4895418040_m01yrx_e5b3a7bd_09.png b/tests/source/foreground/35149c4895418040_m01yrx_e5b3a7bd_09.png
diff --git a/tests/source/foreground/5b9e0751ca458cad_m0gjkl_40736b77_04.png b/tests/source/foreground/5b9e0751ca458cad_m0gjkl_40736b77_04.png
diff --git a/tests/source/foreground/8.jpg b/tests/source/foreground/8.jpg
diff --git a/tests/source/foreground/8351494f79f6bf8f_m03k3r_73176af9_08.png b/tests/source/foreground/8351494f79f6bf8f_m03k3r_73176af9_08.png
diff --git a/tests/source/foreground/f80eda2459853824_m09g1w_b2413ec8_11.png b/tests/source/foreground/f80eda2459853824_m09g1w_b2413ec8_11.png
diff --git a/tests/source/foreground_mask/1.png b/tests/source/foreground_mask/1.png
diff --git a/tests/source/foreground_mask/17.png b/tests/source/foreground_mask/17.png
diff --git a/tests/source/foreground_mask/1bb218825e370c51_m02p5f1q_4a9db762_30.png b/tests/source/foreground_mask/1bb218825e370c51_m02p5f1q_4a9db762_30.png
diff --git a/tests/source/foreground_mask/35149c4895418040_m01yrx_e5b3a7bd_09.png b/tests/source/foreground_mask/35149c4895418040_m01yrx_e5b3a7bd_09.png
diff --git a/tests/source/foreground_mask/5b9e0751ca458cad_m0gjkl_40736b77_04.png b/tests/source/foreground_mask/5b9e0751ca458cad_m0gjkl_40736b77_04.png
diff --git a/tests/source/foreground_mask/8.png b/tests/source/foreground_mask/8.png
diff --git a/tests/source/foreground_mask/8351494f79f6bf8f_m03k3r_73176af9_08.png b/tests/source/foreground_mask/8351494f79f6bf8f_m03k3r_73176af9_08.png
diff --git a/tests/source/foreground_mask/f80eda2459853824_m09g1w_b2413ec8_11.png b/tests/source/foreground_mask/f80eda2459853824_m09g1w_b2413ec8_11.png
diff --git a/tests/test_color_transfer.py b/tests/test_color_transfer.py
@@ -7,33 +7,19 @@
 task_name = 'color_transfer'
 
 if __name__ == '__main__':
-    # test_set = get_test_list()
-    # result_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results', task_name)
-    # if os.path.exists(result_dir):
-    #     shutil.rmtree(result_dir)
-    # os.makedirs(result_dir, exist_ok=True)
-    # print(f'begin testing {task_name}...')
-    # for pair in test_set:
-    #     comp_img, comp_mask = pair['composite'], pair['composite_mask']
-    #     trans_img = color_transfer(comp_img, comp_mask)
-    #     # visulization results
-    #     img_name  = os.path.basename(comp_img).replace('.png', '.jpg')
-    #     grid_img  = make_image_grid([comp_img, comp_mask, trans_img])
-    #     res_path  = os.path.join(result_dir, img_name)
-    #     cv2.imwrite(res_path, grid_img)
-    #     print('save result to ', res_path)
-    # print(f'end testing {task_name}!\n')
-
-    from libcom import color_transfer
-    from libcom.utils.process_image import make_image_grid
-    import cv2
-    comp_img1  = '../tests/source/composite/1.jpg'
-    comp_mask1 = '../tests/source/composite_mask/1.png'
-    trans_img1 = color_transfer(comp_img1, comp_mask1)
-    comp_img2  = '../tests/source/composite/8.jpg'
-    comp_mask2 = '../tests/source/composite_mask/8.png'
-    trans_img2 = color_transfer(comp_img2, comp_mask2)
-    # visualization results
-    grid_img  = make_image_grid([comp_img1, comp_mask1, trans_img1, 
-                                comp_img2, comp_mask2, trans_img2], cols=3)
-    cv2.imwrite('../docs/_static/image/colortransfer_result1.jpg', grid_img)
+    test_set = get_test_list()
+    result_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'results', task_name)
+    if os.path.exists(result_dir):
+        shutil.rmtree(result_dir)
+    os.makedirs(result_dir, exist_ok=True)
+    print(f'begin testing {task_name}...')
+    for pair in test_set:
+        comp_img, comp_mask = pair['composite'], pair['composite_mask']
+        trans_img = color_transfer(comp_img, comp_mask)
+        # visulization results
+        img_name  = os.path.basename(comp_img).replace('.png', '.jpg')
+        grid_img  = make_image_grid([comp_img, comp_mask, trans_img])
+        res_path  = os.path.join(result_dir, img_name)
+        cv2.imwrite(res_path, grid_img)
+        print('save result to ', res_path)
+    print(f'end testing {task_name}!\n')
diff --git a/tests/test_painterly_image_harmonization.py b/tests/test_painterly_image_harmonization.py
@@ -18,8 +18,13 @@
     os.makedirs(result_dir, exist_ok=True)
     print(f'begin testing {task_name}...')
     phd_net = PainterlyHarmonizationModel(device=0, model_type='PHDNet')
-    phdiff  = PainterlyHarmonizationModel(device=0, model_type='PHDiffusion', use_residual=False)
-    phdiffres = PainterlyHarmonizationModel(device=0, model_type='PHDiffusion')
+    gpu_num = torch.cuda.device_count()
+    if gpu_num > 2:
+        phdiff_gpu, phdiffres_gpu = 1, 2
+    elif gpu_num > 1:
+        phdiff_gpu, phdiffres_gpu = 0, 1
+    phdiff  = PainterlyHarmonizationModel(device=phdiff_gpu, model_type='PHDiffusion', use_residual=False)
+    phdiffres = PainterlyHarmonizationModel(device=phdiffres_gpu, model_type='PHDiffusion')
     for pair in test_set:
         comp_img, comp_mask = pair['composite'], pair['composite_mask']
         phdnet_img = phd_net(comp_img, comp_mask)