diff --git a/README.md b/README.md index f47bdb9..aaf25fb 100644 --- a/README.md +++ b/README.md @@ -14,6 +14,13 @@ __Contact__: xuebin[at]ualberta[dot]ca ## Updates !!! +**(2021-Feb-06)** Since some people raised the issue of human segmentation using U^2-Net, we trained another model for human segemntation based on [**Supervisely Person Dataset**](https://supervise.ly/explore/projects/supervisely-person-dataset-23304/datasets). To run the human segmentation model, please first downlowd the [**u2net_human_seg.pth**](https://drive.google.com/file/d/1-Yg0cxgrNhHP-016FPdp902BR-kSsA4P/view?usp=sharing) model weights and start the inference by command: ```python u2net_human_seg_test.py```.
+[**Notes: Due to the labeling accuracy of the Supervisely Person Dataset, the human segmentation model here won't give you hair level accuracy. But it should be more robust than u2net trained with DUTS-TR dataset on human segmentation task. It can be used for human portrait segmentation, human body segmentation, etc.**](https://github.com/NathanUA/U-2-Net)
+ +![Human Image Segmentation](figures/human_seg) +![Human Video](figures/human_seg_video.gif) +![Human Video Results](figures/human_seg_results.gif) + **(2020-Dec-28)** Some interesting applications and useful tools based on U^2-Net:
(1) [**Xiaolong Liu**](https://github.com/LiuXiaolong19920720) developed several very interesting applications based on U^2-Net including [**Human Portrait Drawing**](https://www.cvpy.net/studio/cv/func/DeepLearning/sketch/sketch/page/)(As far as I know, Xiaolong is the first one who uses U^2-Net for portrait generation), [**image matting**](https://www.cvpy.net/studio/cv/func/DeepLearning/matting/matting/page/) and [**so on**](https://www.cvpy.net/).
(2) [**Vladimir Seregin**](https://github.com/peko/nn-lineart) developed an interesting tool, [**NN based lineart**](https://peko.github.io/nn-lineart/), for comparing the portrait results of U^2-Net and that of another popular model, [**ArtLine**](https://github.com/vijishmadhavan/ArtLine), developed by [**Vijish Madhavan**](https://github.com/vijishmadhavan).
diff --git a/__pycache__/data_loader.cpython-38.pyc b/__pycache__/data_loader.cpython-38.pyc new file mode 100644 index 0000000..d4fa72b Binary files /dev/null and b/__pycache__/data_loader.cpython-38.pyc differ diff --git a/figures/human_seg.png b/figures/human_seg.png new file mode 100644 index 0000000..480afc5 Binary files /dev/null and b/figures/human_seg.png differ diff --git a/figures/human_seg_results.gif b/figures/human_seg_results.gif new file mode 100644 index 0000000..93e0c29 Binary files /dev/null and b/figures/human_seg_results.gif differ diff --git a/figures/human_seg_video.gif b/figures/human_seg_video.gif new file mode 100644 index 0000000..d0f542d Binary files /dev/null and b/figures/human_seg_video.gif differ diff --git a/model/__pycache__/__init__.cpython-38.pyc b/model/__pycache__/__init__.cpython-38.pyc new file mode 100644 index 0000000..372489e Binary files /dev/null and b/model/__pycache__/__init__.cpython-38.pyc differ diff --git a/model/__pycache__/u2net.cpython-38.pyc b/model/__pycache__/u2net.cpython-38.pyc new file mode 100644 index 0000000..5a159ec Binary files /dev/null and b/model/__pycache__/u2net.cpython-38.pyc differ diff --git a/test_data/test_human_images/19035828_web1__12294096_web1_180615-PNR-newmayorchallenge.jpg b/test_data/test_human_images/19035828_web1__12294096_web1_180615-PNR-newmayorchallenge.jpg new file mode 100644 index 0000000..236658f Binary files /dev/null and b/test_data/test_human_images/19035828_web1__12294096_web1_180615-PNR-newmayorchallenge.jpg differ diff --git a/test_data/test_human_images/2019-LADIES-NIGHT-2ND-GOMES.jpg b/test_data/test_human_images/2019-LADIES-NIGHT-2ND-GOMES.jpg new file mode 100644 index 0000000..35ff04a Binary files /dev/null and b/test_data/test_human_images/2019-LADIES-NIGHT-2ND-GOMES.jpg differ diff --git a/test_data/test_human_images/480112-637286191695715805-16x9.jpg b/test_data/test_human_images/480112-637286191695715805-16x9.jpg new file mode 100644 index 0000000..966d23f Binary files /dev/null and b/test_data/test_human_images/480112-637286191695715805-16x9.jpg differ diff --git a/test_data/test_human_images/5-mental-skills-of-successful-athletes-image.jpg b/test_data/test_human_images/5-mental-skills-of-successful-athletes-image.jpg new file mode 100644 index 0000000..2959103 Binary files /dev/null and b/test_data/test_human_images/5-mental-skills-of-successful-athletes-image.jpg differ diff --git a/test_data/test_human_images/Athlete-Intake.jpg b/test_data/test_human_images/Athlete-Intake.jpg new file mode 100644 index 0000000..0d0beff Binary files /dev/null and b/test_data/test_human_images/Athlete-Intake.jpg differ diff --git a/test_data/test_human_images/Two_dancers.jpg b/test_data/test_human_images/Two_dancers.jpg new file mode 100644 index 0000000..00fcb1c Binary files /dev/null and b/test_data/test_human_images/Two_dancers.jpg differ diff --git a/test_data/test_human_images/coach-yelling-at-athlete-716268.jpg b/test_data/test_human_images/coach-yelling-at-athlete-716268.jpg new file mode 100644 index 0000000..afa12c6 Binary files /dev/null and b/test_data/test_human_images/coach-yelling-at-athlete-716268.jpg differ diff --git a/test_data/test_human_images/download.jpeg b/test_data/test_human_images/download.jpeg new file mode 100644 index 0000000..894a96e Binary files /dev/null and b/test_data/test_human_images/download.jpeg differ diff --git a/test_data/test_human_images/f2dc6965e31f7ff0bb54618d53437006.jpg b/test_data/test_human_images/f2dc6965e31f7ff0bb54618d53437006.jpg new file mode 100644 index 0000000..9ba10bd Binary files /dev/null and b/test_data/test_human_images/f2dc6965e31f7ff0bb54618d53437006.jpg differ diff --git a/test_data/test_human_images/image1440x560cropped.jpg b/test_data/test_human_images/image1440x560cropped.jpg new file mode 100644 index 0000000..cc4f932 Binary files /dev/null and b/test_data/test_human_images/image1440x560cropped.jpg differ diff --git a/test_data/test_human_images/images (1).jpeg b/test_data/test_human_images/images (1).jpeg new file mode 100644 index 0000000..c669ce3 Binary files /dev/null and b/test_data/test_human_images/images (1).jpeg differ diff --git a/test_data/test_human_images/images (2).jpeg b/test_data/test_human_images/images (2).jpeg new file mode 100644 index 0000000..998edcd Binary files /dev/null and b/test_data/test_human_images/images (2).jpeg differ diff --git a/test_data/test_human_images/images.jpeg b/test_data/test_human_images/images.jpeg new file mode 100644 index 0000000..29c137a Binary files /dev/null and b/test_data/test_human_images/images.jpeg differ diff --git a/test_data/test_human_images/julia+trotti_17.jpg b/test_data/test_human_images/julia+trotti_17.jpg new file mode 100644 index 0000000..47259f2 Binary files /dev/null and b/test_data/test_human_images/julia+trotti_17.jpg differ diff --git a/test_data/test_human_images/language_1280p.jpg b/test_data/test_human_images/language_1280p.jpg new file mode 100644 index 0000000..5e49883 Binary files /dev/null and b/test_data/test_human_images/language_1280p.jpg differ diff --git a/test_data/test_human_images/lionel-messi-athletes-fashion.jpg b/test_data/test_human_images/lionel-messi-athletes-fashion.jpg new file mode 100644 index 0000000..3b6cf07 Binary files /dev/null and b/test_data/test_human_images/lionel-messi-athletes-fashion.jpg differ diff --git a/test_data/test_human_images/olympic-athletes-need-to-know-2018-winter-1.jpg b/test_data/test_human_images/olympic-athletes-need-to-know-2018-winter-1.jpg new file mode 100644 index 0000000..0d54b86 Binary files /dev/null and b/test_data/test_human_images/olympic-athletes-need-to-know-2018-winter-1.jpg differ diff --git a/test_data/test_human_images/photo-1552374196-c4e7ffc6e126.jpeg b/test_data/test_human_images/photo-1552374196-c4e7ffc6e126.jpeg new file mode 100644 index 0000000..ff42890 Binary files /dev/null and b/test_data/test_human_images/photo-1552374196-c4e7ffc6e126.jpeg differ diff --git a/test_data/test_human_images/track-runners-hurdles-1280.jpg b/test_data/test_human_images/track-runners-hurdles-1280.jpg new file mode 100644 index 0000000..88964e7 Binary files /dev/null and b/test_data/test_human_images/track-runners-hurdles-1280.jpg differ diff --git a/test_data/u2net_test_human_images_results/19035828_web1__12294096_web1_180615-PNR-newmayorchallenge.png b/test_data/u2net_test_human_images_results/19035828_web1__12294096_web1_180615-PNR-newmayorchallenge.png new file mode 100644 index 0000000..3e4f5fe Binary files /dev/null and b/test_data/u2net_test_human_images_results/19035828_web1__12294096_web1_180615-PNR-newmayorchallenge.png differ diff --git a/test_data/u2net_test_human_images_results/2019-LADIES-NIGHT-2ND-GOMES.png b/test_data/u2net_test_human_images_results/2019-LADIES-NIGHT-2ND-GOMES.png new file mode 100644 index 0000000..afe16c5 Binary files /dev/null and b/test_data/u2net_test_human_images_results/2019-LADIES-NIGHT-2ND-GOMES.png differ diff --git a/test_data/u2net_test_human_images_results/480112-637286191695715805-16x9.png b/test_data/u2net_test_human_images_results/480112-637286191695715805-16x9.png new file mode 100644 index 0000000..008cce4 Binary files /dev/null and b/test_data/u2net_test_human_images_results/480112-637286191695715805-16x9.png differ diff --git a/test_data/u2net_test_human_images_results/5-mental-skills-of-successful-athletes-image.png b/test_data/u2net_test_human_images_results/5-mental-skills-of-successful-athletes-image.png new file mode 100644 index 0000000..146d29c Binary files /dev/null and b/test_data/u2net_test_human_images_results/5-mental-skills-of-successful-athletes-image.png differ diff --git a/test_data/u2net_test_human_images_results/Athlete-Intake.png b/test_data/u2net_test_human_images_results/Athlete-Intake.png new file mode 100644 index 0000000..d75aa6d Binary files /dev/null and b/test_data/u2net_test_human_images_results/Athlete-Intake.png differ diff --git a/test_data/u2net_test_human_images_results/Two_dancers.png b/test_data/u2net_test_human_images_results/Two_dancers.png new file mode 100644 index 0000000..bbb1bc9 Binary files /dev/null and b/test_data/u2net_test_human_images_results/Two_dancers.png differ diff --git a/test_data/u2net_test_human_images_results/coach-yelling-at-athlete-716268.png b/test_data/u2net_test_human_images_results/coach-yelling-at-athlete-716268.png new file mode 100644 index 0000000..4b977ef Binary files /dev/null and b/test_data/u2net_test_human_images_results/coach-yelling-at-athlete-716268.png differ diff --git a/test_data/u2net_test_human_images_results/download.png b/test_data/u2net_test_human_images_results/download.png new file mode 100644 index 0000000..750d494 Binary files /dev/null and b/test_data/u2net_test_human_images_results/download.png differ diff --git a/test_data/u2net_test_human_images_results/f2dc6965e31f7ff0bb54618d53437006.png b/test_data/u2net_test_human_images_results/f2dc6965e31f7ff0bb54618d53437006.png new file mode 100644 index 0000000..0681400 Binary files /dev/null and b/test_data/u2net_test_human_images_results/f2dc6965e31f7ff0bb54618d53437006.png differ diff --git a/test_data/u2net_test_human_images_results/image1440x560cropped.png b/test_data/u2net_test_human_images_results/image1440x560cropped.png new file mode 100644 index 0000000..0fec915 Binary files /dev/null and b/test_data/u2net_test_human_images_results/image1440x560cropped.png differ diff --git a/test_data/u2net_test_human_images_results/images (1).png b/test_data/u2net_test_human_images_results/images (1).png new file mode 100644 index 0000000..9f32688 Binary files /dev/null and b/test_data/u2net_test_human_images_results/images (1).png differ diff --git a/test_data/u2net_test_human_images_results/images (2).png b/test_data/u2net_test_human_images_results/images (2).png new file mode 100644 index 0000000..a3372a7 Binary files /dev/null and b/test_data/u2net_test_human_images_results/images (2).png differ diff --git a/test_data/u2net_test_human_images_results/images.png b/test_data/u2net_test_human_images_results/images.png new file mode 100644 index 0000000..f0b1886 Binary files /dev/null and b/test_data/u2net_test_human_images_results/images.png differ diff --git a/test_data/u2net_test_human_images_results/julia+trotti_17.png b/test_data/u2net_test_human_images_results/julia+trotti_17.png new file mode 100644 index 0000000..8e69584 Binary files /dev/null and b/test_data/u2net_test_human_images_results/julia+trotti_17.png differ diff --git a/test_data/u2net_test_human_images_results/language_1280p.png b/test_data/u2net_test_human_images_results/language_1280p.png new file mode 100644 index 0000000..aefa791 Binary files /dev/null and b/test_data/u2net_test_human_images_results/language_1280p.png differ diff --git a/test_data/u2net_test_human_images_results/lionel-messi-athletes-fashion.png b/test_data/u2net_test_human_images_results/lionel-messi-athletes-fashion.png new file mode 100644 index 0000000..7ba1f6b Binary files /dev/null and b/test_data/u2net_test_human_images_results/lionel-messi-athletes-fashion.png differ diff --git a/test_data/u2net_test_human_images_results/olympic-athletes-need-to-know-2018-winter-1.png b/test_data/u2net_test_human_images_results/olympic-athletes-need-to-know-2018-winter-1.png new file mode 100644 index 0000000..99a4ae1 Binary files /dev/null and b/test_data/u2net_test_human_images_results/olympic-athletes-need-to-know-2018-winter-1.png differ diff --git a/test_data/u2net_test_human_images_results/photo-1552374196-c4e7ffc6e126.png b/test_data/u2net_test_human_images_results/photo-1552374196-c4e7ffc6e126.png new file mode 100644 index 0000000..f4de465 Binary files /dev/null and b/test_data/u2net_test_human_images_results/photo-1552374196-c4e7ffc6e126.png differ diff --git a/test_data/u2net_test_human_images_results/track-runners-hurdles-1280.png b/test_data/u2net_test_human_images_results/track-runners-hurdles-1280.png new file mode 100644 index 0000000..3c0ee0e Binary files /dev/null and b/test_data/u2net_test_human_images_results/track-runners-hurdles-1280.png differ diff --git a/u2net_human_seg_test.py b/u2net_human_seg_test.py new file mode 100644 index 0000000..53c5e35 --- /dev/null +++ b/u2net_human_seg_test.py @@ -0,0 +1,117 @@ +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 + +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 + +# normalize the predicted SOD probability map +def normPRED(d): + ma = torch.max(d) + mi = torch.min(d) + + dn = (d-mi)/(ma-mi) + + 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) + + 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') + +def main(): + + # --------- 1. get image path and name --------- + model_name='u2net' + + + image_dir = os.path.join(os.getcwd(), 'test_data', 'test_human_images')#'test_human_images')#'test_portrait_images', 'your_portrait_im') + prediction_dir = os.path.join(os.getcwd(), 'test_data', 'test_human_images' + '_results' + os.sep) + model_dir = os.path.join(os.getcwd(), 'saved_models', model_name+'_human_seg', model_name + '_human_seg.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 = [], + transform=transforms.Compose([RescaleT(320), + ToTensorLab(flag=0)]) + ) + test_salobj_dataloader = DataLoader(test_salobj_dataset, + batch_size=1, + shuffle=False, + num_workers=1) + + # --------- 3. model define --------- + if(model_name=='u2net'): + print("...load U2NET---173.6 MB") + net = U2NET(3,1) + + if torch.cuda.is_available(): + net.load_state_dict(torch.load(model_dir)) + net.cuda() + else: + net.load_state_dict(torch.load(model_dir, map_location='cpu')) + net.eval() + + # --------- 4. inference for each image --------- + for i_test, data_test in enumerate(test_salobj_dataloader): + + print("inferencing:",img_name_list[i_test].split(os.sep)[-1]) + + inputs_test = data_test['image'] + inputs_test = inputs_test.type(torch.FloatTensor) + + if torch.cuda.is_available(): + inputs_test = Variable(inputs_test.cuda()) + else: + inputs_test = Variable(inputs_test) + + d1,d2,d3,d4,d5,d6,d7= net(inputs_test) + + # normalization + pred = d1[:,0,:,:] + pred = normPRED(pred) + + # 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) + + del d1,d2,d3,d4,d5,d6,d7 + +if __name__ == "__main__": + main()