diff --git a/README.md b/README.md index 9573517..b972b3c 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,18 @@ __Contact__: xuebin[at]ualberta[dot]ca ## Updates !!! +**(2021-Aug-24)** We play a bit more about fusing the orignal image and the generated portraits to composite different styles. You can (1) download this repo by: +``` +git clone https://github.com/NathanUA/U-2-Net.git +``` +, (2) Download the u2net_portrait.pth from [**GoogleDrive**](https://drive.google.com/file/d/1IG3HdpcRiDoWNookbncQjeaPN28t90yW/view?usp=sharing) or [**Baidu Pan(提取码：chgd)**](https://pan.baidu.com/s/1BYT5Ts6BxwpB8_l2sAyCkw)model and put it into the directory: ```./saved_models/u2net_portrait/```, (3) run the code by command +``` +python u2net_portrait_composite.py -s 20 -a 0.5 +``` +,where ``-s`` indicates the sigmoid of gaussian function for blurring the orignal image and ``-a`` denotes the alpha weights of the orignal image when fusing them. + +![im_portrait_composite](figures/im_composite.jpg') + **(2021-July-16)** A new [background removal webapp](https://remove-background.net/) developed by Изатоп Василий. ![rm_bg](figures/rm_bg.JPG) diff --git a/__pycache__/data_loader.cpython-38.pyc b/__pycache__/data_loader.cpython-38.pyc index d4fa72b..2cd2a60 100644 Binary files a/__pycache__/data_loader.cpython-38.pyc and b/__pycache__/data_loader.cpython-38.pyc differ diff --git a/figures/im_composite.jpg b/figures/im_composite.jpg new file mode 100644 index 0000000..ff5f596 Binary files /dev/null and b/figures/im_composite.jpg differ diff --git a/model/__pycache__/__init__.cpython-38.pyc b/model/__pycache__/__init__.cpython-38.pyc index 372489e..e243a8a 100644 Binary files a/model/__pycache__/__init__.cpython-38.pyc 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 index 5a159ec..2cc6fd8 100644 Binary files a/model/__pycache__/u2net.cpython-38.pyc and b/model/__pycache__/u2net.cpython-38.pyc differ diff --git a/test_data/test_portrait_images/portrait_results/img_1695.png b/test_data/test_portrait_images/portrait_results/img_1695.png index 770bc36..9c5ebaa 100644 Binary files a/test_data/test_portrait_images/portrait_results/img_1695.png and b/test_data/test_portrait_images/portrait_results/img_1695.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/GalGadot_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/GalGadot_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..2583776 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/GalGadot_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/guliNazha3_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/guliNazha3_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..6bdae30 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/guliNazha3_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/kid1_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/kid1_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..461fa9f Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/kid1_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/kid2_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/kid2_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..80c2a06 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/kid2_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/kid3_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/kid3_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..39bf8f8 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/kid3_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/man2_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/man2_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..2d3e320 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/man2_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/man4_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/man4_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..f1a4787 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/man4_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/man5_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/man5_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..3df8294 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/man5_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/man_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/man_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..88796e3 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/man_sigma_20.0_alpha_0.5_composite.png differ diff --git a/test_data/test_portrait_images/your_portrait_results/smile_sigma_20.0_alpha_0.5_composite.png b/test_data/test_portrait_images/your_portrait_results/smile_sigma_20.0_alpha_0.5_composite.png new file mode 100644 index 0000000..103d808 Binary files /dev/null and b/test_data/test_portrait_images/your_portrait_results/smile_sigma_20.0_alpha_0.5_composite.png differ diff --git a/u2net_portrait_composite.py b/u2net_portrait_composite.py new file mode 100644 index 0000000..74d5da5 --- /dev/null +++ b/u2net_portrait_composite.py @@ -0,0 +1,141 @@ +import os +from skimage import io, transform +from skimage.filters import gaussian +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 +from model import U2NETP # small version u2net 4.7 MB + +import argparse + +# 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,sigma=2,alpha=0.5): + + predict = pred + predict = predict.squeeze() + predict_np = predict.cpu().data.numpy() + + image = io.imread(image_name) + pd = transform.resize(predict_np,image.shape[0:2],order=2) + pd = pd/(np.amax(pd)+1e-8)*255 + pd = pd[:,:,np.newaxis] + + print(image.shape) + print(pd.shape) + + ## fuse the orignal portrait image and the portraits into one composite image + ## 1. use gaussian filter to blur the orginal image + sigma=sigma + image = gaussian(image, sigma=sigma, preserve_range=True) + + ## 2. fuse these orignal image and the portrait with certain weight: alpha + alpha = alpha + im_comp = image*alpha+pd*(1-alpha) + + print(im_comp.shape) + + + img_name = image_name.split(os.sep)[-1] + aaa = img_name.split(".") + bbb = aaa[0:-1] + imidx = bbb[0] + for i in range(1,len(bbb)): + imidx = imidx + "." + bbb[i] + io.imsave(d_dir+'/'+imidx+'_sigma_' + str(sigma) + '_alpha_' + str(alpha) + '_composite.png',im_comp) + +def main(): + + parser = argparse.ArgumentParser(description="image and portrait composite") + parser.add_argument('-s',action='store',dest='sigma') + parser.add_argument('-a',action='store',dest='alpha') + args = parser.parse_args() + print(args.sigma) + print(args.alpha) + print("--------------------") + + # --------- 1. get image path and name --------- + model_name='u2net_portrait'#u2netp + + + image_dir = './test_data/test_portrait_images/your_portrait_im' + prediction_dir = './test_data/test_portrait_images/your_portrait_results' + if(not os.path.exists(prediction_dir)): + os.mkdir(prediction_dir) + + model_dir = './saved_models/u2net_portrait/u2net_portrait.pth' + + img_name_list = glob.glob(image_dir+'/*') + print("Number of images: ", len(img_name_list)) + + # --------- 2. dataloader --------- + #1. dataloader + test_salobj_dataset = SalObjDataset(img_name_list = img_name_list, + lbl_name_list = [], + transform=transforms.Compose([RescaleT(512), + ToTensorLab(flag=0)]) + ) + test_salobj_dataloader = DataLoader(test_salobj_dataset, + batch_size=1, + shuffle=False, + num_workers=1) + + # --------- 3. model define --------- + + print("...load U2NET---173.6 MB") + net = U2NET(3,1) + + net.load_state_dict(torch.load(model_dir)) + if torch.cuda.is_available(): + net.cuda() + 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 = 1.0 - d1[:,0,:,:] + pred = normPRED(pred) + + # save results to test_results folder + save_output(img_name_list[i_test],pred,prediction_dir,sigma=float(args.sigma),alpha=float(args.alpha)) + + del d1,d2,d3,d4,d5,d6,d7 + +if __name__ == "__main__": + main()