xuebinqin · LIUZHENYU0609 · Jul 26, 2021 · Jul 26, 2021 · Aug 2, 2021
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,144 @@
+### Python template
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+saved_models/u2net
+saved_models/u2netp
+.idea
+test_data
diff --git a/.idea/.gitignore b/.idea/.gitignore
diff --git a/.idea/U-2-Net.iml b/.idea/U-2-Net.iml
diff --git a/.idea/inspectionProfiles/Project_Default.xml b/.idea/inspectionProfiles/Project_Default.xml
diff --git a/.idea/inspectionProfiles/profiles_settings.xml b/.idea/inspectionProfiles/profiles_settings.xml
diff --git a/.idea/misc.xml b/.idea/misc.xml
diff --git a/.idea/modules.xml b/.idea/modules.xml
diff --git a/.idea/vcs.xml b/.idea/vcs.xml
diff --git a/model/__pycache__/__init__.cpython-36.pyc b/model/__pycache__/__init__.cpython-36.pyc
diff --git a/model/__pycache__/u2net.cpython-36.pyc b/model/__pycache__/u2net.cpython-36.pyc
diff --git a/test_data/test_images/0002-01.jpg b/test_data/test_images/0002-01.jpg
diff --git a/test_data/test_images/0003.jpg b/test_data/test_images/0003.jpg
diff --git a/test_data/test_images/bike.jpg b/test_data/test_images/bike.jpg
diff --git a/test_data/test_images/boat.jpg b/test_data/test_images/boat.jpg
diff --git a/test_data/test_images/girl.png b/test_data/test_images/girl.png
diff --git a/test_data/test_images/hockey.png b/test_data/test_images/hockey.png
diff --git a/test_data/test_images/horse.jpg b/test_data/test_images/horse.jpg
diff --git a/test_data/test_images/im_01.png b/test_data/test_images/im_01.png
diff --git a/test_data/test_images/im_14.png b/test_data/test_images/im_14.png
diff --git a/test_data/test_images/im_21.png b/test_data/test_images/im_21.png
diff --git a/test_data/test_images/im_27.png b/test_data/test_images/im_27.png
diff --git a/test_data/test_images/lamp2_meitu_1.jpg b/test_data/test_images/lamp2_meitu_1.jpg
diff --git a/test_data/test_images/long.jpg b/test_data/test_images/long.jpg
diff --git a/test_data/test_images/rifle1.jpg b/test_data/test_images/rifle1.jpg
diff --git a/test_data/test_images/rifle2.jpeg b/test_data/test_images/rifle2.jpeg
diff --git a/test_data/test_images/sailboat3.jpg b/test_data/test_images/sailboat3.jpg
diff --git a/test_data/test_images/vangogh.jpeg b/test_data/test_images/vangogh.jpeg
diff --git a/test_data/test_images/whisk.png b/test_data/test_images/whisk.png
diff --git a/test_data/u2net_results/0002-01.png b/test_data/u2net_results/0002-01.png
diff --git a/test_data/u2net_results/0003.png b/test_data/u2net_results/0003.png
diff --git a/test_data/u2net_results/bike.png b/test_data/u2net_results/bike.png
diff --git a/test_data/u2net_results/boat.png b/test_data/u2net_results/boat.png
diff --git a/test_data/u2net_results/girl.png b/test_data/u2net_results/girl.png
diff --git a/test_data/u2net_results/hockey.png b/test_data/u2net_results/hockey.png
diff --git a/test_data/u2net_results/horse.png b/test_data/u2net_results/horse.png
diff --git a/test_data/u2net_results/im_01.png b/test_data/u2net_results/im_01.png
diff --git a/test_data/u2net_results/im_14.png b/test_data/u2net_results/im_14.png
diff --git a/test_data/u2net_results/im_21.png b/test_data/u2net_results/im_21.png
diff --git a/test_data/u2net_results/im_27.png b/test_data/u2net_results/im_27.png
diff --git a/test_data/u2net_results/lamp2_meitu_1.png b/test_data/u2net_results/lamp2_meitu_1.png
diff --git a/test_data/u2net_results/long.png b/test_data/u2net_results/long.png
diff --git a/test_data/u2net_results/rifle1.png b/test_data/u2net_results/rifle1.png
diff --git a/test_data/u2net_results/rifle2.png b/test_data/u2net_results/rifle2.png
diff --git a/test_data/u2net_results/sailboat3.png b/test_data/u2net_results/sailboat3.png
diff --git a/test_data/u2net_results/vangogh.png b/test_data/u2net_results/vangogh.png
diff --git a/test_data/u2net_results/whisk.png b/test_data/u2net_results/whisk.png
diff --git a/test_data/u2netp_results/0002-01.png b/test_data/u2netp_results/0002-01.png
diff --git a/test_data/u2netp_results/0003.png b/test_data/u2netp_results/0003.png
diff --git a/test_data/u2netp_results/bike.png b/test_data/u2netp_results/bike.png
diff --git a/test_data/u2netp_results/boat.png b/test_data/u2netp_results/boat.png
diff --git a/test_data/u2netp_results/girl.png b/test_data/u2netp_results/girl.png
diff --git a/test_data/u2netp_results/hockey.png b/test_data/u2netp_results/hockey.png
diff --git a/test_data/u2netp_results/horse.png b/test_data/u2netp_results/horse.png
diff --git a/test_data/u2netp_results/im_01.png b/test_data/u2netp_results/im_01.png
diff --git a/test_data/u2netp_results/im_14.png b/test_data/u2netp_results/im_14.png
diff --git a/test_data/u2netp_results/im_21.png b/test_data/u2netp_results/im_21.png
diff --git a/test_data/u2netp_results/im_27.png b/test_data/u2netp_results/im_27.png
diff --git a/test_data/u2netp_results/lamp2_meitu_1.png b/test_data/u2netp_results/lamp2_meitu_1.png
diff --git a/test_data/u2netp_results/long.png b/test_data/u2netp_results/long.png
diff --git a/test_data/u2netp_results/rifle1.png b/test_data/u2netp_results/rifle1.png
diff --git a/test_data/u2netp_results/rifle2.png b/test_data/u2netp_results/rifle2.png
diff --git a/test_data/u2netp_results/sailboat3.png b/test_data/u2netp_results/sailboat3.png
diff --git a/test_data/u2netp_results/vangogh.png b/test_data/u2netp_results/vangogh.png
diff --git a/test_data/u2netp_results/whisk.png b/test_data/u2netp_results/whisk.png
diff --git a/u2net_test.py b/u2net_test.py
@@ -1,25 +1,22 @@
 import os
 from skimage import io, transform
 import torch
-import torchvision
 from torch.autograd import Variable
-import torch.nn as nn
-import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
-from torchvision import transforms#, utils
-# import torch.optim as optim
+from torchvision import transforms
 
 import numpy as np
 from PIL import Image
 import glob
 
 from data_loader import RescaleT
-from data_loader import ToTensor
 from data_loader import ToTensorLab
 from data_loader import SalObjDataset
 
-from model import U2NET # full size version 173.6 MB
-from model import U2NETP # small version u2net 4.7 MB
+from model import U2NET
+from model import U2NETP
+import cv2
+
 
 # normalize the predicted SOD probability map
 def normPRED(d):
@@ -30,45 +27,63 @@ def normPRED(d):
 
     return dn
 
+
 def save_output(image_name,pred,d_dir):
 
     predict = pred
     predict = predict.squeeze()
     predict_np = predict.cpu().data.numpy()
-
     im = Image.fromarray(predict_np*255).convert('RGB')
     img_name = image_name.split(os.sep)[-1]
     image = io.imread(image_name)
-    imo = im.resize((image.shape[1],image.shape[0]),resample=Image.BILINEAR)
-
-    pb_np = np.array(imo)
+    imo = im.resize((image.shape[1], image.shape[0]), resample=Image.BILINEAR)
 
     aaa = img_name.split(".")
     bbb = aaa[0:-1]
     imidx = bbb[0]
     for i in range(1,len(bbb)):
         imidx = imidx + "." + bbb[i]
 
-    imo.save(d_dir+imidx+'.png')
+    mask_save_dir = os.path.join(d_dir, "mask_results"+os.sep)
+    if not os.path.exists(mask_save_dir):
+        os.makedirs(mask_save_dir, exist_ok=True)
+    imo.save(mask_save_dir+imidx+'.png')
 
-def main():
 
-    # --------- 1. get image path and name ---------
-    model_name='u2net'#u2netp
+# 根据mask输出获取前置图像
+def save_front_image(image_name, prediction_dir):
+    img1 = cv2.imread(image_name)
+    img_name = image_name.split(os.sep)[-1]
+    indexes = img_name.split(".")[0:-1]
+    imidx = indexes[0]
+    for i in range(1, len(indexes)):
+        imidx = imidx + "." + indexes[i]
+    mask_image_path = os.path.join(prediction_dir, "mask_results"+os.sep, imidx+'.png')
+    img2 = cv2.imread(mask_image_path, cv2.IMREAD_GRAYSCALE)
+    h, w, c = img1.shape
+    img3 = np.zeros((h, w, 4))
+    img3[:, :, 0:3] = img1
+    img3[:, :, 3] = img2
+    front_save_dir = os.path.join(prediction_dir, "front_results"+os.sep)
+    if not os.path.exists(front_save_dir):
+        os.makedirs(front_save_dir, exist_ok=True)
+    cv2.imwrite(front_save_dir+imidx+'.png', img3)
 
 
+def main():
 
+    # --------- 1. get image path and name ---------
+    model_name='u2net'#u2netp
     image_dir = os.path.join(os.getcwd(), 'test_data', 'test_images')
     prediction_dir = os.path.join(os.getcwd(), 'test_data', model_name + '_results' + os.sep)
     model_dir = os.path.join(os.getcwd(), 'saved_models', model_name, model_name + '.pth')
-
     img_name_list = glob.glob(image_dir + os.sep + '*')
     print(img_name_list)
 
     # --------- 2. dataloader ---------
-    #1. dataloader
-    test_salobj_dataset = SalObjDataset(img_name_list = img_name_list,
-                                        lbl_name_list = [],
+    # 1. dataloader
+    test_salobj_dataset = SalObjDataset(img_name_list=img_name_list,
+                                        lbl_name_list=[],
                                         transform=transforms.Compose([RescaleT(320),
                                                                       ToTensorLab(flag=0)])
                                         )
@@ -114,9 +129,11 @@ def main():
         # save results to test_results folder
         if not os.path.exists(prediction_dir):
             os.makedirs(prediction_dir, exist_ok=True)
-        save_output(img_name_list[i_test],pred,prediction_dir)
+        save_output(img_name_list[i_test], pred, prediction_dir)
+        save_front_image(img_name_list[i_test], prediction_dir)
+
+        del d1, d2, d3, d4, d5, d6, d7
 
-        del d1,d2,d3,d4,d5,d6,d7
 
 if __name__ == "__main__":
     main()