使用U2net+cpu+pytorch完成照片素描化

u2net_portrait_test.py

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
from model import U2NETP # small version u2net 4.7 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_portrait'#u2netp

image_dir = './test_data/test_portrait_images/portrait_im'
prediction_dir = './test_data/test_portrait_images/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('./saved_models/u2net_portrait/u2net_portrait.pth',map_location='cpu'))
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)

del d1,d2,d3,d4,d5,d6,d7

if __name__ == "__main__":
main()

使用U2net+cpu+pytorch完成照片素描化相关推荐

php图片素描化,照片变素描软件,只需一次点,照片变素描|图片素描
我们在大街小巷人人热闹的地方经常会看到街头有人在给人素描,那么我们在电脑上能用照片素描转换软件来直接把照片转换成素描画吗?答案是可以的,今天咱们分享给大家的就是照片素描转换软件,可以一键将照片转换成素 ...
Facebook将照片3D化技术商用，强大的算法+海量的数据+移动端优化是工程亮点
点击我爱计算机视觉标星,更快获取CVML新技术近日,Facebook AI 团队官宣使用其 Facebook APP的用户均可以方便的在手机端将单摄像头拍摄的2D图像转成3D图像,这项服务支持iPh ...
win7小工具打不开_强大图片素描化工具 FotoSketcher 3.4
软件介绍 FotoSketcher 是一款很简便且强大的图片素描化工具,处理速度很快,可以帮助你创建类似真实素描的艺术作品. 更新内容/软件特色 FotoSketcher 现在有二十多种自动素描方式, ...
如何在ps中将图片素描化
如何在photoshop中将图片素描化作者:陈涛写作时间:2019.1.17 其实 ,在Photoshop在将一张图片转变成素描图片非常的简单,接下来我就为大家讲解一下如何将图片素描化首先,我们 ...
ps小知识——将图片素描化
....... (PS小知识--将图片素描化) 步骤: 一:先复制背景图层快捷键:Ctrl + J 二:点击"图像"-"调整"-"去色" ...
代码实现照片素描_我的代码素描之旅
代码实现照片素描 In 2018 I started the process of consistently creating and posting my code sketches online. ...
python 素描化图像_任意图像转素描：Python分分钟实现
图像转素描效果展示看到这个标题你可能觉得,我这次终于开始标题党了.然而众所周知,我是标图党~,一般不会做标题党这种事情,既然说了分分钟实现,那就说明--代码在十行左右. 代码虽少,但是效果确实一点都 ...
php图片素描化,把照片做成素描效果照片做成素描
专注实拍监控风景各类视频申请入群加微信: neimu500 因为专注所以专业! 我们只发优质实拍视频,群管理专业.规范,确保群的视频质量以及秩序!无水印可下载可分享. 网上流传有很多使用ps ...
Sketch n Cartoonize for mac(照片卡通化素描化编辑工具)
Sketch n Cartoonize 可以使用复杂的算法来实现"逼真的"和"近乎完美"的卡通外观,获得您想要的照片效果. Sketch n Cartooniz ...

使用U2net+cpu+pytorch完成照片素描化

使用U2net+cpu+pytorch完成照片素描化相关推荐

最新文章

热门文章