提交 d2444b7e 编写于 作者: W wuzewu

Add thirdparty directory

上级 7cad8e4e
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self,in_ch=3,out_ch=3,dirate=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2D(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer):#UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer):#UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer):#UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer):#UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer):#UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx2d = self.rebnconv2d(paddle.concat((hx3d,hx2),1))
hx1d = self.rebnconv1d(paddle.concat((hx2d,hx1),1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NET,self).__init__()
self.stage1 = RSU7(in_ch,32,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(128,out_ch,3,padding=1)
self.side4 = nn.Conv2D(256,out_ch,3,padding=1)
self.side5 = nn.Conv2D(512,out_ch,3,padding=1)
self.side6 = nn.Conv2D(512,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NETP,self).__init__()
self.stage1 = RSU7(in_ch,16,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,16,64)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(64,16,64)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(64,16,64)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(64,16,64)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(64,16,64)
# decoder
self.stage5d = RSU4F(128,16,64)
self.stage4d = RSU4(128,16,64)
self.stage3d = RSU5(128,16,64)
self.stage2d = RSU6(128,16,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(64,out_ch,3,padding=1)
self.side4 = nn.Conv2D(64,out_ch,3,padding=1)
self.side5 = nn.Conv2D(64,out_ch,3,padding=1)
self.side6 = nn.Conv2D(64,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self,in_ch=3,out_ch=3,dirate=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2D(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer):#UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer):#UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer):#UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer):#UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer):#UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx2d = self.rebnconv2d(paddle.concat((hx3d,hx2),1))
hx1d = self.rebnconv1d(paddle.concat((hx2d,hx1),1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NET,self).__init__()
self.stage1 = RSU7(in_ch,32,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(128,out_ch,3,padding=1)
self.side4 = nn.Conv2D(256,out_ch,3,padding=1)
self.side5 = nn.Conv2D(512,out_ch,3,padding=1)
self.side6 = nn.Conv2D(512,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NETP,self).__init__()
self.stage1 = RSU7(in_ch,16,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,16,64)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(64,16,64)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(64,16,64)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(64,16,64)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(64,16,64)
# decoder
self.stage5d = RSU4F(128,16,64)
self.stage4d = RSU4(128,16,64)
self.stage3d = RSU5(128,16,64)
self.stage2d = RSU6(128,16,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(64,out_ch,3,padding=1)
self.side4 = nn.Conv2D(64,out_ch,3,padding=1)
self.side5 = nn.Conv2D(64,out_ch,3,padding=1)
self.side6 = nn.Conv2D(64,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self,in_ch=3,out_ch=3,dirate=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2D(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer):#UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer):#UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer):#UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer):#UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer):#UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx2d = self.rebnconv2d(paddle.concat((hx3d,hx2),1))
hx1d = self.rebnconv1d(paddle.concat((hx2d,hx1),1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NET,self).__init__()
self.stage1 = RSU7(in_ch,32,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(128,out_ch,3,padding=1)
self.side4 = nn.Conv2D(256,out_ch,3,padding=1)
self.side5 = nn.Conv2D(512,out_ch,3,padding=1)
self.side6 = nn.Conv2D(512,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NETP,self).__init__()
self.stage1 = RSU7(in_ch,16,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,16,64)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(64,16,64)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(64,16,64)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(64,16,64)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(64,16,64)
# decoder
self.stage5d = RSU4F(128,16,64)
self.stage4d = RSU4(128,16,64)
self.stage3d = RSU5(128,16,64)
self.stage2d = RSU6(128,16,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(64,out_ch,3,padding=1)
self.side4 = nn.Conv2D(64,out_ch,3,padding=1)
self.side5 = nn.Conv2D(64,out_ch,3,padding=1)
self.side6 = nn.Conv2D(64,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
......@@ -7,13 +7,14 @@ import paddle.nn as nn
import paddlehub as hub
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="ID_Photo_GEN", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="ID_Photo_GEN", # 模型介绍
version="1.0.0" # 版本号
name="ID_Photo_GEN", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="ID_Photo_GEN", # 模型介绍
version="1.0.0" # 版本号
)
class ID_Photo_GEN(nn.Layer):
def __init__(self):
......@@ -38,18 +39,15 @@ class ID_Photo_GEN(nn.Layer):
if images is not None:
datas = images
# 返回数据列表
return datas
# 数据预处理函数
def preprocess(self, images, batch_size, use_gpu):
# 获取人脸关键点
outputs = self.face_detector.keypoint_detection(
images=images,
batch_size=batch_size,
use_gpu=use_gpu)
outputs = self.face_detector.keypoint_detection(images=images, batch_size=batch_size, use_gpu=use_gpu)
crops = []
for output, image in zip(outputs, images):
for landmarks in output['data']:
......@@ -58,7 +56,8 @@ class ID_Photo_GEN(nn.Layer):
# rotation angle
left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45]
radian = np.arctan((left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
radian = np.arctan(
(left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
# image size after rotating
height, width, _ = image.shape
......@@ -73,7 +72,7 @@ class ID_Photo_GEN(nn.Layer):
# affine matrix
M = np.array([[cos, sin, (1 - cos) * width / 2. - sin * height / 2. + Tx],
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
image = cv2.warpAffine(image, M, (new_w, new_h), borderValue=(255, 255, 255))
......@@ -103,16 +102,17 @@ class ID_Photo_GEN(nn.Layer):
h, w = image.shape[:2]
left_white = max(0, -left)
left = max(0, left)
right = min(right, w-1)
right_white = left_white + (right-left)
right = min(right, w - 1)
right_white = left_white + (right - left)
top_white = max(0, -top)
top = max(0, top)
bottom = min(bottom, h-1)
bottom = min(bottom, h - 1)
bottom_white = top_white + (bottom - top)
image_crop[top_white:bottom_white+1, left_white:right_white+1] = image[top:bottom+1, left:right+1].copy()
image_crop[top_white:bottom_white + 1, left_white:right_white + 1] = image[top:bottom + 1, left:right +
1].copy()
crops.append(image_crop)
# 获取人像分割的输出
results = self.seg.Segmentation(images=crops, batch_size=batch_size)
......@@ -141,7 +141,7 @@ class ID_Photo_GEN(nn.Layer):
cartoon = self.net(data)
outputs.append(cartoon[0].numpy())
outputs = np.concatenate(outputs, 0)
return outputs
......@@ -168,22 +168,11 @@ class ID_Photo_GEN(nn.Layer):
cv2.imwrite(os.path.join(output_dir, 'blue_%d.jpg' % i), blue)
cv2.imwrite(os.path.join(output_dir, 'red_%d.jpg' % i), red)
results.append({
'write': write,
'blue': blue,
'red': red
})
results.append({'write': write, 'blue': blue, 'red': red})
return results
def Photo_GEN(
self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
def Photo_GEN(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False, use_gpu=False):
# 获取输入数据
images = self.load_datas(paths, images)
......@@ -194,4 +183,4 @@ class ID_Photo_GEN(nn.Layer):
# 结果后处理
results = self.postprocess(faces, masks, visualization, output_dir)
return results
\ No newline at end of file
return results
......@@ -17,72 +17,66 @@ class ResnetGenerator(nn.Layer):
nn.ReLU()
]
DownBlock += [
HourGlass(ngf, ngf),
HourGlass(ngf, ngf)
]
DownBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf)]
# Down-Sampling
n_downsampling = 2
for i in range(n_downsampling):
mult = 2 ** i
mult = 2**i
DownBlock += [
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf*mult, ngf*mult*2, kernel_size=3, stride=2, bias_attr=False),
nn.InstanceNorm2D(ngf*mult*2, weight_attr=False, bias_attr=False),
nn.Conv2D(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, bias_attr=False),
nn.InstanceNorm2D(ngf * mult * 2, weight_attr=False, bias_attr=False),
nn.ReLU()
]
# Encoder Bottleneck
mult = 2 ** n_downsampling
mult = 2**n_downsampling
for i in range(n_blocks):
setattr(self, 'EncodeBlock'+str(i+1), ResnetBlock(ngf*mult))
setattr(self, 'EncodeBlock' + str(i + 1), ResnetBlock(ngf * mult))
# Class Activation Map
self.gap_fc = nn.Linear(ngf*mult, 1, bias_attr=False)
self.gmp_fc = nn.Linear(ngf*mult, 1, bias_attr=False)
self.conv1x1 = nn.Conv2D(ngf*mult*2, ngf*mult, kernel_size=1, stride=1)
self.gap_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.gmp_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.conv1x1 = nn.Conv2D(ngf * mult * 2, ngf * mult, kernel_size=1, stride=1)
self.relu = nn.ReLU()
# Gamma, Beta block
FC = []
if self.light:
FC += [
nn.Linear(ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(),
nn.Linear(ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
else:
FC += [
nn.Linear(img_size//mult*img_size//mult*ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(img_size // mult * img_size // mult * ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(),
nn.Linear(ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
# Decoder Bottleneck
mult = 2 ** n_downsampling
mult = 2**n_downsampling
for i in range(n_blocks):
setattr(self, 'DecodeBlock'+str(i + 1), ResnetSoftAdaLINBlock(ngf*mult))
setattr(self, 'DecodeBlock' + str(i + 1), ResnetSoftAdaLINBlock(ngf * mult))
# Up-Sampling
UpBlock = []
for i in range(n_downsampling):
mult = 2 ** (n_downsampling - i)
mult = 2**(n_downsampling - i)
UpBlock += [
nn.Upsample(scale_factor=2),
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf*mult, ngf*mult//2, kernel_size=3, stride=1, bias_attr=False),
LIN(ngf*mult//2),
nn.Conv2D(ngf * mult, ngf * mult // 2, kernel_size=3, stride=1, bias_attr=False),
LIN(ngf * mult // 2),
nn.ReLU()
]
UpBlock += [
HourGlass(ngf, ngf),
HourGlass(ngf, ngf, False)
]
UpBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf, False)]
UpBlock += [
nn.Pad2D([3, 3, 3, 3], 'reflect'),
......@@ -101,7 +95,7 @@ class ResnetGenerator(nn.Layer):
content_features = []
for i in range(self.n_blocks):
x = getattr(self, 'EncodeBlock'+str(i+1))(x)
x = getattr(self, 'EncodeBlock' + str(i + 1))(x)
content_features.append(F.adaptive_avg_pool2d(x, 1).reshape([bs, -1]))
gap = F.adaptive_avg_pool2d(x, 1)
......@@ -127,7 +121,7 @@ class ResnetGenerator(nn.Layer):
style_features = self.FC(x.reshape([bs, -1]))
for i in range(self.n_blocks):
x = getattr(self, 'DecodeBlock'+str(i+1))(x, content_features[4-i-1], style_features)
x = getattr(self, 'DecodeBlock' + str(i + 1))(x, content_features[4 - i - 1], style_features)
out = self.UpBlock(x)
......@@ -140,25 +134,20 @@ class ConvBlock(nn.Layer):
self.dim_in = dim_in
self.dim_out = dim_out
self.conv_block1 = self.__convblock(dim_in, dim_out//2)
self.conv_block2 = self.__convblock(dim_out//2, dim_out//4)
self.conv_block3 = self.__convblock(dim_out//4, dim_out//4)
self.conv_block1 = self.__convblock(dim_in, dim_out // 2)
self.conv_block2 = self.__convblock(dim_out // 2, dim_out // 4)
self.conv_block3 = self.__convblock(dim_out // 4, dim_out // 4)
if self.dim_in != self.dim_out:
self.conv_skip = nn.Sequential(
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False),
nn.ReLU(),
nn.Conv2D(dim_in, dim_out, kernel_size=1, stride=1, bias_attr=False)
)
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.ReLU(),
nn.Conv2D(dim_in, dim_out, kernel_size=1, stride=1, bias_attr=False))
@staticmethod
def __convblock(dim_in, dim_out):
return nn.Sequential(
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False),
nn.ReLU(),
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(dim_in, dim_out, kernel_size=3, stride=1, bias_attr=False)
)
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.ReLU(), nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(dim_in, dim_out, kernel_size=3, stride=1, bias_attr=False))
def forward(self, x):
residual = x
......@@ -182,24 +171,24 @@ class HourGlassBlock(nn.Layer):
self.n_block = 9
for i in range(self.n_skip):
setattr(self, 'ConvBlockskip'+str(i+1), ConvBlock(dim_in, dim_in))
setattr(self, 'ConvBlockskip' + str(i + 1), ConvBlock(dim_in, dim_in))
for i in range(self.n_block):
setattr(self, 'ConvBlock'+str(i+1), ConvBlock(dim_in, dim_in))
setattr(self, 'ConvBlock' + str(i + 1), ConvBlock(dim_in, dim_in))
def forward(self, x):
skips = []
for i in range(self.n_skip):
skips.append(getattr(self, 'ConvBlockskip'+str(i+1))(x))
skips.append(getattr(self, 'ConvBlockskip' + str(i + 1))(x))
x = F.avg_pool2d(x, 2)
x = getattr(self, 'ConvBlock'+str(i+1))(x)
x = getattr(self, 'ConvBlock' + str(i + 1))(x)
x = self.ConvBlock5(x)
for i in range(self.n_skip):
x = getattr(self, 'ConvBlock'+str(i+6))(x)
x = getattr(self, 'ConvBlock' + str(i + 6))(x)
x = F.upsample(x, scale_factor=2)
x = skips[self.n_skip-i-1] + x
x = skips[self.n_skip - i - 1] + x
return x
......@@ -210,12 +199,9 @@ class HourGlass(nn.Layer):
self.use_res = use_res
self.HG = nn.Sequential(
HourGlassBlock(dim_in),
ConvBlock(dim_out, dim_out),
HourGlassBlock(dim_in), ConvBlock(dim_out, dim_out),
nn.Conv2D(dim_out, dim_out, kernel_size=1, stride=1, bias_attr=False),
nn.InstanceNorm2D(dim_out, weight_attr=False, bias_attr=False),
nn.ReLU()
)
nn.InstanceNorm2D(dim_out, weight_attr=False, bias_attr=False), nn.ReLU())
self.Conv1 = nn.Conv2D(dim_out, 3, kernel_size=1, stride=1)
......@@ -292,12 +278,12 @@ class SoftAdaLIN(nn.Layer):
self.w_gamma = self.create_parameter([1, num_features], default_initializer=nn.initializer.Constant(0.))
self.w_beta = self.create_parameter([1, num_features], default_initializer=nn.initializer.Constant(0.))
self.c_gamma = nn.Sequential(nn.Linear(num_features, num_features, bias_attr=False),
nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.c_beta = nn.Sequential(nn.Linear(num_features, num_features, bias_attr=False),
nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.c_gamma = nn.Sequential(
nn.Linear(num_features, num_features, bias_attr=False), nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.c_beta = nn.Sequential(
nn.Linear(num_features, num_features, bias_attr=False), nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.s_gamma = nn.Linear(num_features, num_features, bias_attr=False)
self.s_beta = nn.Linear(num_features, num_features, bias_attr=False)
......@@ -352,8 +338,8 @@ class LIN(nn.Layer):
out = out * self.gamma.expand([x.shape[0], -1, -1, -1]) + self.beta.expand([x.shape[0], -1, -1, -1])
return out
if __name__ == '__main__':
#d = Discriminator(3)
# paddle.summary(d, (4, 3, 256, 256))
......
......@@ -8,13 +8,14 @@ import paddlehub as hub
from Photo2Cartoon.model import ResnetGenerator
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="Photo2Cartoon", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="Photo2Cartoon", # 模型介绍
version="1.0.0" # 版本号
name="Photo2Cartoon", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="Photo2Cartoon", # 模型介绍
version="1.0.0" # 版本号
)
class Photo2Cartoon(nn.Layer):
def __init__(self):
......@@ -49,18 +50,15 @@ class Photo2Cartoon(nn.Layer):
if images is not None:
datas = images
# 返回数据列表
return datas
# 数据预处理函数
def preprocess(self, images, batch_size, use_gpu):
# 获取人脸关键点
outputs = self.face_detector.keypoint_detection(
images=images,
batch_size=batch_size,
use_gpu=use_gpu)
outputs = self.face_detector.keypoint_detection(images=images, batch_size=batch_size, use_gpu=use_gpu)
crops = []
for output, image in zip(outputs, images):
for landmarks in output['data']:
......@@ -69,7 +67,8 @@ class Photo2Cartoon(nn.Layer):
# rotation angle
left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45]
radian = np.arctan((left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
radian = np.arctan(
(left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
# image size after rotating
height, width, _ = image.shape
......@@ -84,7 +83,7 @@ class Photo2Cartoon(nn.Layer):
# affine matrix
M = np.array([[cos, sin, (1 - cos) * width / 2. - sin * height / 2. + Tx],
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
image = cv2.warpAffine(image, M, (new_w, new_h), borderValue=(255, 255, 255))
......@@ -114,16 +113,17 @@ class Photo2Cartoon(nn.Layer):
h, w = image.shape[:2]
left_white = max(0, -left)
left = max(0, left)
right = min(right, w-1)
right_white = left_white + (right-left)
right = min(right, w - 1)
right_white = left_white + (right - left)
top_white = max(0, -top)
top = max(0, top)
bottom = min(bottom, h-1)
bottom = min(bottom, h - 1)
bottom_white = top_white + (bottom - top)
image_crop[top_white:bottom_white+1, left_white:right_white+1] = image[top:bottom+1, left:right+1].copy()
image_crop[top_white:bottom_white + 1, left_white:right_white + 1] = image[top:bottom + 1, left:right +
1].copy()
crops.append(image_crop)
# 获取人像分割的输出
results = self.seg.Segmentation(images=crops, batch_size=batch_size)
......@@ -158,7 +158,7 @@ class Photo2Cartoon(nn.Layer):
# 切分数据
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
return input_datas, masks
......@@ -175,7 +175,7 @@ class Photo2Cartoon(nn.Layer):
cartoon = self.net(data)
outputs.append(cartoon[0].numpy())
outputs = np.concatenate(outputs, 0)
return outputs
......@@ -204,17 +204,16 @@ class Photo2Cartoon(nn.Layer):
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), cartoon)
cartoons.append(cartoon)
return cartoons
def Cartoon_GEN(
self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
def Cartoon_GEN(self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
# 获取输入数据
images = self.load_datas(paths, images)
......@@ -228,4 +227,4 @@ class Photo2Cartoon(nn.Layer):
# 结果后处理
cartoons = self.postprocess(outputs, masks, visualization, output_dir)
return cartoons
\ No newline at end of file
return cartoons
......@@ -6,18 +6,19 @@ from U2Net_Portrait.u2net import U2NET
from U2Net_Portrait.processor import Processor
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="U2Net_Portrait", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Net_Portrait", # 模型介绍
version="1.0.0" # 版本号
name="U2Net_Portrait", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Net_Portrait", # 模型介绍
version="1.0.0" # 版本号
)
class U2Net_Portrait(nn.Layer):
def __init__(self):
super(U2Net_Portrait, self).__init__()
self.model = U2NET(3,1)
self.model = U2NET(3, 1)
state_dict = paddle.load(os.path.join(self.directory, 'u2net_portrait.pdparams'))
self.model.set_dict(state_dict)
self.model.eval()
......@@ -26,22 +27,21 @@ class U2Net_Portrait(nn.Layer):
outputs = []
for data in input_datas:
data = paddle.to_tensor(data, 'float32')
d1,d2,d3,d4,d5,d6,d7= self.model(data)
d1, d2, d3, d4, d5, d6, d7 = self.model(data)
outputs.append(d1.numpy())
outputs = np.concatenate(outputs, 0)
return outputs
def Portrait_GEN(
self,
images=None,
paths=None,
scale=1,
batch_size=1,
output_dir='output',
face_detection=True,
visualization=False):
def Portrait_GEN(self,
images=None,
paths=None,
scale=1,
batch_size=1,
output_dir='output',
face_detection=True,
visualization=False):
# 初始化数据处理器
processor = Processor(paths, images, batch_size, face_detection, scale)
......
......@@ -5,6 +5,7 @@ import paddlehub as hub
__all__ = ['Processor']
class Processor():
def __init__(self, paths, images, batch_size, face_detection=True, scale=1):
# 图像列表
......@@ -26,7 +27,7 @@ class Processor():
if images is not None:
datas = images
# 返回数据列表
return datas
......@@ -35,10 +36,7 @@ class Processor():
if face_detection:
# face detection
face_detector = hub.Module(name="pyramidbox_lite_mobile")
results = face_detector.face_detection(images=imgs,
use_gpu=False,
visualization=False,
confs_threshold=0.5)
results = face_detector.face_detection(images=imgs, use_gpu=False, visualization=False, confs_threshold=0.5)
im_faces = []
for datas, img in zip(results, imgs):
for face in datas['data']:
......@@ -46,52 +44,64 @@ class Processor():
l, r, t, b = [face['left'], face['right'], face['top'], face['bottom']]
# square crop
pad = max(int(scale*(r-l)), int(scale*(b-t)))
c_w, c_h = (r-l)//2+l, (b-t)//2+t
top = 0 if c_h-pad<0 else c_h-pad
pad = max(int(scale * (r - l)), int(scale * (b - t)))
c_w, c_h = (r - l) // 2 + l, (b - t) // 2 + t
top = 0 if c_h - pad < 0 else c_h - pad
bottom = pad + c_h
left = 0 if c_w-pad<0 else c_w-pad
left = 0 if c_w - pad < 0 else c_w - pad
right = pad + c_w
crop = img[top:bottom, left:right]
# resize
im_face = cv2.resize(crop, (512,512), interpolation = cv2.INTER_AREA)
im_face = cv2.resize(crop, (512, 512), interpolation=cv2.INTER_AREA)
im_faces.append(im_face)
else:
im_faces = []
for img in imgs:
h, w = img.shape[:2]
if h>w:
if (h-w)%2==0:
img = np.pad(img,((0,0),((h-w)//2,(h-w)//2),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
if h > w:
if (h - w) % 2 == 0:
img = np.pad(
img, ((0, 0), ((h - w) // 2, (h - w) // 2), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
else:
img = np.pad(img,((0,0),((h-w)//2,(h-w)//2+1),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
img = np.pad(
img, ((0, 0), ((h - w) // 2, (h - w) // 2 + 1), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
else:
if (w-h)%2==0:
img = np.pad(img,(((w-h)//2,(w-h)//2),(0,0),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
if (w - h) % 2 == 0:
img = np.pad(
img, (((w - h) // 2, (w - h) // 2), (0, 0), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
else:
img = np.pad(img,(((w-h)//2,(w-h)//2+1),(0,0),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
im_face = cv2.resize(img, (512,512), interpolation = cv2.INTER_AREA)
img = np.pad(
img, (((w - h) // 2, (w - h) // 2 + 1), (0, 0), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
im_face = cv2.resize(img, (512, 512), interpolation=cv2.INTER_AREA)
im_faces.append(im_face)
input_datas = []
for im_face in im_faces:
tmpImg = np.zeros((im_face.shape[0],im_face.shape[1],3))
im_face = im_face/np.max(im_face)
tmpImg = np.zeros((im_face.shape[0], im_face.shape[1], 3))
im_face = im_face / np.max(im_face)
tmpImg[:,:,0] = (im_face[:,:,2]-0.406)/0.225
tmpImg[:,:,1] = (im_face[:,:,1]-0.456)/0.224
tmpImg[:,:,2] = (im_face[:,:,0]-0.485)/0.229
tmpImg[:, :, 0] = (im_face[:, :, 2] - 0.406) / 0.225
tmpImg[:, :, 1] = (im_face[:, :, 1] - 0.456) / 0.224
tmpImg[:, :, 2] = (im_face[:, :, 0] - 0.485) / 0.229
# convert BGR to RGB
tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = tmpImg[np.newaxis,:,:,:]
tmpImg = tmpImg[np.newaxis, :, :, :]
input_datas.append(tmpImg)
input_datas = np.concatenate(input_datas, 0)
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
......@@ -101,7 +111,7 @@ class Processor():
ma = np.max(d)
mi = np.min(d)
dn = (d-mi)/(ma-mi)
dn = (d - mi) / (ma - mi)
return dn
......@@ -110,20 +120,20 @@ class Processor():
results = []
if visualization and not os.path.exists(output_dir):
os.mkdir(output_dir)
for i in range(outputs.shape[0]):
# normalization
pred = 1.0 - outputs[i,0,:,:]
pred = 1.0 - outputs[i, 0, :, :]
pred = self.normPRED(pred)
# convert torch tensor to numpy array
pred = pred.squeeze()
pred = (pred*255).astype(np.uint8)
pred = (pred * 255).astype(np.uint8)
results.append(pred)
if visualization:
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), pred)
return results
\ No newline at end of file
return results
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self, in_ch=3, out_ch=3, dirate=1):
super(REBNCONV, self).__init__()
self.conv_s1 = nn.Conv2D(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self, x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src, tar):
src = F.upsample(src, size=tar.shape[2:], mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer): #UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool5 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7, hx6), 1))
hx6dup = _upsample_like(hx6d, hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer): #UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer): #UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer): #UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer): #UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx2d = self.rebnconv2d(paddle.concat((hx3d, hx2), 1))
hx1d = self.rebnconv1d(paddle.concat((hx2d, hx1), 1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NET, self).__init__()
self.stage1 = RSU7(in_ch, 32, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 32, 128)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(128, 64, 256)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(256, 128, 512)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(512, 256, 512)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(512, 256, 512)
# decoder
self.stage5d = RSU4F(1024, 256, 512)
self.stage4d = RSU4(1024, 128, 256)
self.stage3d = RSU5(512, 64, 128)
self.stage2d = RSU6(256, 32, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(128, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(256, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(512, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(512, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NETP, self).__init__()
self.stage1 = RSU7(in_ch, 16, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 16, 64)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(64, 16, 64)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(64, 16, 64)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(64, 16, 64)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(64, 16, 64)
# decoder
self.stage5d = RSU4F(128, 16, 64)
self.stage4d = RSU4(128, 16, 64)
self.stage3d = RSU5(128, 16, 64)
self.stage2d = RSU6(128, 16, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(64, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['InferenceModel']
class InferenceModel():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True):
......@@ -25,14 +26,10 @@ class InferenceModel():
# 打印函数
def __repr__(self):
'''
get the numbers and name of inputs and outputs
get the numbers and name of inputs and outputs
'''
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (
len(self.input_handles),
str(self.input_names),
len(self.output_handles),
str(self.output_names)
)
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (len(
self.input_handles), str(self.input_names), len(self.output_handles), str(self.output_names))
# 类调用函数
def __call__(self, *input_datas, batch_size=1):
......@@ -59,9 +56,11 @@ class InferenceModel():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -71,7 +70,7 @@ class InferenceModel():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -125,17 +124,17 @@ class InferenceModel():
self.input_handles[i].copy_from_cpu(input_data)
self.predictor.run()
for i in range(len(self.output_handles)):
output = self.output_handles[i].copy_to_cpu()
if i in outputs:
outputs[i].append(output)
else:
outputs[i] = [output]
# 预测结果合并
for key in outputs.keys():
outputs[key] = np.concatenate(outputs[key], 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -26,12 +26,7 @@ class MiDaS_Large(Module):
model_path = os.path.join(self.directory, "model-f6b98070")
# 加载模型
self.model = InferenceModel(
modelpath=model_path,
use_gpu=use_gpu,
use_mkldnn=False,
combined=True
)
self.model = InferenceModel(modelpath=model_path, use_gpu=use_gpu, use_mkldnn=False, combined=True)
self.model.eval()
# 数据预处理配置
......@@ -46,11 +41,10 @@ class MiDaS_Large(Module):
resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC,
),
NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
PrepareForNet()
])
# 数据读取函数
@staticmethod
def load_datas(paths, images):
......@@ -84,7 +78,7 @@ class MiDaS_Large(Module):
input_data = img[np.newaxis, ...]
input_datas.append(input_data)
# 拼接数据
input_datas = np.concatenate(input_datas, 0)
......@@ -113,12 +107,7 @@ class MiDaS_Large(Module):
return outputs
# 深度估计函数
def depth_estimation(self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
def depth_estimation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
# 加载数据
datas = self.load_datas(paths, images)
......@@ -130,5 +119,5 @@ class MiDaS_Large(Module):
# 结果后处理
outputs = self.postprocess(datas, results, output_dir, visualization)
return outputs
\ No newline at end of file
return outputs
......@@ -7,6 +7,7 @@ import cv2
class Resize(object):
"""Resize sample to given size (width, height).
"""
def __init__(self,
width,
height,
......@@ -51,12 +52,10 @@ class Resize(object):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
if y < min_val:
y = (np.ceil(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
return y
......@@ -91,31 +90,24 @@ class Resize(object):
# fit height
scale_width = scale_height
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
if self.__resize_method == "lower_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
elif self.__resize_method == "upper_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
elif self.__resize_method == "minimal":
new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width)
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
return (new_width, new_height)
def __call__(self, sample):
width, height = self.get_size(sample["image"].shape[1],
sample["image"].shape[0])
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
# resize sample
sample["image"] = cv2.resize(
......@@ -133,8 +125,7 @@ class Resize(object):
)
if "depth" in sample:
sample["depth"] = cv2.resize(sample["depth"], (width, height),
interpolation=cv2.INTER_NEAREST)
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
sample["mask"] = cv2.resize(
sample["mask"].astype(np.float32),
......@@ -149,6 +140,7 @@ class Resize(object):
class NormalizeImage(object):
"""Normlize image by given mean and std.
"""
def __init__(self, mean, std):
self.__mean = mean
self.__std = std
......@@ -162,6 +154,7 @@ class NormalizeImage(object):
class PrepareForNet(object):
"""Prepare sample for usage as network input.
"""
def __init__(self):
pass
......
......@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (len(image.shape) == 2
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
elif (len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False
else:
raise Exception(
"Image must have H x W x 3, H x W x 1 or H x W dimensions.")
raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
file.write("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
......
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['InferenceModel']
class InferenceModel():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True):
......@@ -25,14 +26,10 @@ class InferenceModel():
# 打印函数
def __repr__(self):
'''
get the numbers and name of inputs and outputs
get the numbers and name of inputs and outputs
'''
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (
len(self.input_handles),
str(self.input_names),
len(self.output_handles),
str(self.output_names)
)
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (len(
self.input_handles), str(self.input_names), len(self.output_handles), str(self.output_names))
# 类调用函数
def __call__(self, *input_datas, batch_size=1):
......@@ -59,9 +56,11 @@ class InferenceModel():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -71,7 +70,7 @@ class InferenceModel():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -125,17 +124,17 @@ class InferenceModel():
self.input_handles[i].copy_from_cpu(input_data)
self.predictor.run()
for i in range(len(self.output_handles)):
output = self.output_handles[i].copy_to_cpu()
if i in outputs:
outputs[i].append(output)
else:
outputs[i] = [output]
# 预测结果合并
for key in outputs.keys():
outputs[key] = np.concatenate(outputs[key], 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -26,12 +26,7 @@ class MiDaS_Small(Module):
model_path = os.path.join(self.directory, "model-small")
# 加载模型
self.model = InferenceModel(
modelpath=model_path,
use_gpu=use_gpu,
use_mkldnn=False,
combined=True
)
self.model = InferenceModel(modelpath=model_path, use_gpu=use_gpu, use_mkldnn=False, combined=True)
self.model.eval()
# 数据预处理配置
......@@ -46,11 +41,10 @@ class MiDaS_Small(Module):
resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC,
),
NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
PrepareForNet()
])
# 数据读取函数
@staticmethod
def load_datas(paths, images):
......@@ -84,7 +78,7 @@ class MiDaS_Small(Module):
input_data = img[np.newaxis, ...]
input_datas.append(input_data)
# 拼接数据
input_datas = np.concatenate(input_datas, 0)
......@@ -113,12 +107,7 @@ class MiDaS_Small(Module):
return outputs
# 深度估计函数
def depth_estimation(self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
def depth_estimation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
# 加载数据
datas = self.load_datas(paths, images)
......@@ -130,5 +119,5 @@ class MiDaS_Small(Module):
# 结果后处理
outputs = self.postprocess(datas, results, output_dir, visualization)
return outputs
\ No newline at end of file
return outputs
......@@ -7,6 +7,7 @@ import cv2
class Resize(object):
"""Resize sample to given size (width, height).
"""
def __init__(self,
width,
height,
......@@ -51,12 +52,10 @@ class Resize(object):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
if y < min_val:
y = (np.ceil(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
return y
......@@ -91,31 +90,24 @@ class Resize(object):
# fit height
scale_width = scale_height
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
if self.__resize_method == "lower_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
elif self.__resize_method == "upper_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
elif self.__resize_method == "minimal":
new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width)
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
return (new_width, new_height)
def __call__(self, sample):
width, height = self.get_size(sample["image"].shape[1],
sample["image"].shape[0])
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
# resize sample
sample["image"] = cv2.resize(
......@@ -133,8 +125,7 @@ class Resize(object):
)
if "depth" in sample:
sample["depth"] = cv2.resize(sample["depth"], (width, height),
interpolation=cv2.INTER_NEAREST)
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
sample["mask"] = cv2.resize(
sample["mask"].astype(np.float32),
......@@ -149,6 +140,7 @@ class Resize(object):
class NormalizeImage(object):
"""Normlize image by given mean and std.
"""
def __init__(self, mean, std):
self.__mean = mean
self.__std = std
......@@ -162,6 +154,7 @@ class NormalizeImage(object):
class PrepareForNet(object):
"""Prepare sample for usage as network input.
"""
def __init__(self):
pass
......
......@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (len(image.shape) == 2
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
elif (len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False
else:
raise Exception(
"Image must have H x W x 3, H x W x 1 or H x W dimensions.")
raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
file.write("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
......
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
......@@ -2,10 +2,11 @@ import numpy as np
import cv2
from scipy import ndimage
def get_normal_map(img):
img = img.astype(np.float)
img = img / 255.0
img = - img + 1
img = -img + 1
img[img < 0] = 0
img[img > 1] = 1
return img
......@@ -112,7 +113,7 @@ def resize_img_512(img):
def resize_img_512_3d(img):
zeros = np.zeros((1, 3, 512, 512), dtype=np.float)
zeros[0, 0: img.shape[0], 0: img.shape[1], 0: img.shape[2]] = img
zeros[0, 0:img.shape[0], 0:img.shape[1], 0:img.shape[2]] = img
return zeros.transpose((1, 2, 3, 0))
......@@ -122,7 +123,7 @@ def denoise_mat(img, i):
def show_active_img_and_save_denoise(img, path):
mat = img.astype(np.float)
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......@@ -132,11 +133,9 @@ def show_active_img_and_save_denoise(img, path):
return
def show_active_img(name, img):
mat = img.astype(np.float)
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......@@ -147,7 +146,7 @@ def show_active_img(name, img):
def get_active_img(img):
mat = img.astype(np.float)
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......@@ -158,7 +157,7 @@ def get_active_img(img):
def get_active_img_fil(img):
mat = img.astype(np.float)
mat[mat < 0.18] = 0
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......
......@@ -14,12 +14,12 @@ import paddle.fluid as fluid
import paddlehub as hub
from Extract_Line_Draft.function import *
@moduleinfo(
name="Extract_Line_Draft",
version="1.0.0",
type="cv/segmentation",
summary=
"Import the color picture and generate the line draft of the picture",
summary="Import the color picture and generate the line draft of the picture",
author="彭兆帅,郑博培",
author_email="1084667371@qq.com,2733821739@qq.com")
class ExtractLineDraft(hub.Module):
......@@ -28,9 +28,10 @@ class ExtractLineDraft(hub.Module):
Initialize with the necessary elements
"""
# 加载模型路径
self.default_pretrained_model_path = os.path.join(self.directory, "assets","infer_model")
self.default_pretrained_model_path = os.path.join(self.directory, "assets", "infer_model")
self._set_config()
def _set_config(self):
def _set_config(self):
"""
predictor config setting
"""
......@@ -72,7 +73,7 @@ class ExtractLineDraft(hub.Module):
self.predictor.zero_copy_run()
output = self.output_tensor.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
......@@ -153,12 +154,9 @@ class ExtractLineDraft(hub.Module):
usage='%(prog)s',
add_help=True)
self.arg_input_group = self.parser.add_argument_group(
title="Input options", description="Input data. Required")
self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required")
self.arg_config_group = self.parser.add_argument_group(
title="Config options",
description=
"Run configuration for controlling module behavior, not required.")
title="Config options", description="Run configuration for controlling module behavior, not required.")
self.add_module_input_arg()
......@@ -173,21 +171,12 @@ class ExtractLineDraft(hub.Module):
use_gpu = args.use_gpu
self.ExtractLine(image=input_data, use_gpu=use_gpu)
def add_module_input_arg(self):
"""
Add the command input options
"""
self.arg_input_group.add_argument(
'--image',
type=str,
default=None,
help="file contain input data")
self.arg_input_group.add_argument(
'--use_gpu',
type=ast.literal_eval,
default=None,
help="weather to use gpu")
self.arg_input_group.add_argument('--image', type=str, default=None, help="file contain input data")
self.arg_input_group.add_argument('--use_gpu', type=ast.literal_eval, default=None, help="weather to use gpu")
def check_input_data(self, args):
input_data = []
......@@ -196,4 +185,3 @@ class ExtractLineDraft(hub.Module):
raise RuntimeError("Path %s is not exist." % args.image)
path = "{}".format(args.image)
return path
......@@ -48,12 +48,9 @@ class FCN(nn.Layer):
super(FCN, self).__init__()
self.backbone = backbone
backbone_channels = [
backbone.feat_channels[i] for i in backbone_indices
]
backbone_channels = [backbone.feat_channels[i] for i in backbone_indices]
self.head = FCNHead(num_classes, backbone_indices, backbone_channels,
channels)
self.head = FCNHead(num_classes, backbone_indices, backbone_channels, channels)
self.align_corners = align_corners
self.pretrained = pretrained
......@@ -62,11 +59,7 @@ class FCN(nn.Layer):
feat_list = self.backbone(x)
logit_list = self.head(feat_list)
return [
F.interpolate(
logit,
x.shape[2:],
mode='bilinear',
align_corners=self.align_corners) for logit in logit_list
F.interpolate(logit, x.shape[2:], mode='bilinear', align_corners=self.align_corners) for logit in logit_list
]
......@@ -83,11 +76,7 @@ class FCNHead(nn.Layer):
pretrained (str, optional): The path of pretrained model. Default: None
"""
def __init__(self,
num_classes,
backbone_indices=(-1, ),
backbone_channels=(270, ),
channels=None):
def __init__(self, num_classes, backbone_indices=(-1, ), backbone_channels=(270, ), channels=None):
super(FCNHead, self).__init__()
self.num_classes = num_classes
......@@ -96,17 +85,8 @@ class FCNHead(nn.Layer):
channels = backbone_channels[0]
self.conv_1 = ConvBNReLU(
in_channels=backbone_channels[0],
out_channels=channels,
kernel_size=1,
padding='same',
stride=1)
self.cls = nn.Conv2D(
in_channels=channels,
out_channels=self.num_classes,
kernel_size=1,
stride=1,
padding=0)
in_channels=backbone_channels[0], out_channels=channels, kernel_size=1, padding='same', stride=1)
self.cls = nn.Conv2D(in_channels=channels, out_channels=self.num_classes, kernel_size=1, stride=1, padding=0)
def forward(self, feat_list):
logit_list = []
......@@ -115,4 +95,3 @@ class FCNHead(nn.Layer):
logit = self.cls(x)
logit_list.append(logit)
return logit_list
......@@ -21,8 +21,8 @@ import paddle.nn.functional as F
from .layers import ConvBNReLU, ConvBN
__all__ = [
"HRNet_W18_Small_V1", "HRNet_W18_Small_V2", "HRNet_W18", "HRNet_W30",
"HRNet_W32", "HRNet_W40", "HRNet_W44", "HRNet_W48", "HRNet_W60", "HRNet_W64"
"HRNet_W18_Small_V1", "HRNet_W18_Small_V2", "HRNet_W18", "HRNet_W30", "HRNet_W32", "HRNet_W40", "HRNet_W44",
"HRNet_W48", "HRNet_W60", "HRNet_W64"
]
......@@ -88,20 +88,10 @@ class HRNet(nn.Layer):
self.feat_channels = [sum(stage4_num_channels)]
self.conv_layer1_1 = ConvBNReLU(
in_channels=3,
out_channels=64,
kernel_size=3,
stride=2,
padding='same',
bias_attr=False)
in_channels=3, out_channels=64, kernel_size=3, stride=2, padding='same', bias_attr=False)
self.conv_layer1_2 = ConvBNReLU(
in_channels=64,
out_channels=64,
kernel_size=3,
stride=2,
padding='same',
bias_attr=False)
in_channels=64, out_channels=64, kernel_size=3, stride=2, padding='same', bias_attr=False)
self.la1 = Layer1(
num_channels=64,
......@@ -111,9 +101,7 @@ class HRNet(nn.Layer):
name="layer2")
self.tr1 = TransitionLayer(
in_channels=[self.stage1_num_channels[0] * 4],
out_channels=self.stage2_num_channels,
name="tr1")
in_channels=[self.stage1_num_channels[0] * 4], out_channels=self.stage2_num_channels, name="tr1")
self.st2 = Stage(
num_channels=self.stage2_num_channels,
......@@ -125,9 +113,7 @@ class HRNet(nn.Layer):
align_corners=align_corners)
self.tr2 = TransitionLayer(
in_channels=self.stage2_num_channels,
out_channels=self.stage3_num_channels,
name="tr2")
in_channels=self.stage2_num_channels, out_channels=self.stage3_num_channels, name="tr2")
self.st3 = Stage(
num_channels=self.stage3_num_channels,
num_modules=self.stage3_num_modules,
......@@ -138,9 +124,7 @@ class HRNet(nn.Layer):
align_corners=align_corners)
self.tr3 = TransitionLayer(
in_channels=self.stage3_num_channels,
out_channels=self.stage4_num_channels,
name="tr3")
in_channels=self.stage3_num_channels, out_channels=self.stage4_num_channels, name="tr3")
self.st4 = Stage(
num_channels=self.stage4_num_channels,
num_modules=self.stage4_num_modules,
......@@ -166,30 +150,16 @@ class HRNet(nn.Layer):
st4 = self.st4(tr3)
x0_h, x0_w = st4[0].shape[2:]
x1 = F.interpolate(
st4[1], (x0_h, x0_w),
mode='bilinear',
align_corners=self.align_corners)
x2 = F.interpolate(
st4[2], (x0_h, x0_w),
mode='bilinear',
align_corners=self.align_corners)
x3 = F.interpolate(
st4[3], (x0_h, x0_w),
mode='bilinear',
align_corners=self.align_corners)
x1 = F.interpolate(st4[1], (x0_h, x0_w), mode='bilinear', align_corners=self.align_corners)
x2 = F.interpolate(st4[2], (x0_h, x0_w), mode='bilinear', align_corners=self.align_corners)
x3 = F.interpolate(st4[3], (x0_h, x0_w), mode='bilinear', align_corners=self.align_corners)
x = paddle.concat([st4[0], x1, x2, x3], axis=1)
return [x]
class Layer1(nn.Layer):
def __init__(self,
num_channels,
num_filters,
num_blocks,
has_se=False,
name=None):
def __init__(self, num_channels, num_filters, num_blocks, has_se=False, name=None):
super(Layer1, self).__init__()
self.bottleneck_block_list = []
......@@ -258,12 +228,7 @@ class TransitionLayer(nn.Layer):
class Branches(nn.Layer):
def __init__(self,
num_blocks,
in_channels,
out_channels,
has_se=False,
name=None):
def __init__(self, num_blocks, in_channels, out_channels, has_se=False, name=None):
super(Branches, self).__init__()
self.basic_block_list = []
......@@ -278,8 +243,7 @@ class Branches(nn.Layer):
num_channels=in_ch,
num_filters=out_channels[i],
has_se=has_se,
name=name + '_branch_layer_' + str(i + 1) + '_' +
str(j + 1)))
name=name + '_branch_layer_' + str(i + 1) + '_' + str(j + 1)))
self.basic_block_list[i].append(basic_block_func)
def forward(self, x):
......@@ -293,24 +257,14 @@ class Branches(nn.Layer):
class BottleneckBlock(nn.Layer):
def __init__(self,
num_channels,
num_filters,
has_se,
stride=1,
downsample=False,
name=None):
def __init__(self, num_channels, num_filters, has_se, stride=1, downsample=False, name=None):
super(BottleneckBlock, self).__init__()
self.has_se = has_se
self.downsample = downsample
self.conv1 = ConvBNReLU(
in_channels=num_channels,
out_channels=num_filters,
kernel_size=1,
padding='same',
bias_attr=False)
in_channels=num_channels, out_channels=num_filters, kernel_size=1, padding='same', bias_attr=False)
self.conv2 = ConvBNReLU(
in_channels=num_filters,
......@@ -321,26 +275,15 @@ class BottleneckBlock(nn.Layer):
bias_attr=False)
self.conv3 = ConvBN(
in_channels=num_filters,
out_channels=num_filters * 4,
kernel_size=1,
padding='same',
bias_attr=False)
in_channels=num_filters, out_channels=num_filters * 4, kernel_size=1, padding='same', bias_attr=False)
if self.downsample:
self.conv_down = ConvBN(
in_channels=num_channels,
out_channels=num_filters * 4,
kernel_size=1,
padding='same',
bias_attr=False)
in_channels=num_channels, out_channels=num_filters * 4, kernel_size=1, padding='same', bias_attr=False)
if self.has_se:
self.se = SELayer(
num_channels=num_filters * 4,
num_filters=num_filters * 4,
reduction_ratio=16,
name=name + '_fc')
num_channels=num_filters * 4, num_filters=num_filters * 4, reduction_ratio=16, name=name + '_fc')
def forward(self, x):
residual = x
......@@ -360,13 +303,7 @@ class BottleneckBlock(nn.Layer):
class BasicBlock(nn.Layer):
def __init__(self,
num_channels,
num_filters,
stride=1,
has_se=False,
downsample=False,
name=None):
def __init__(self, num_channels, num_filters, stride=1, has_se=False, downsample=False, name=None):
super(BasicBlock, self).__init__()
self.has_se = has_se
......@@ -380,26 +317,14 @@ class BasicBlock(nn.Layer):
padding='same',
bias_attr=False)
self.conv2 = ConvBN(
in_channels=num_filters,
out_channels=num_filters,
kernel_size=3,
padding='same',
bias_attr=False)
in_channels=num_filters, out_channels=num_filters, kernel_size=3, padding='same', bias_attr=False)
if self.downsample:
self.conv_down = ConvBNReLU(
in_channels=num_channels,
out_channels=num_filters,
kernel_size=1,
padding='same',
bias_attr=False)
in_channels=num_channels, out_channels=num_filters, kernel_size=1, padding='same', bias_attr=False)
if self.has_se:
self.se = SELayer(
num_channels=num_filters,
num_filters=num_filters,
reduction_ratio=16,
name=name + '_fc')
self.se = SELayer(num_channels=num_filters, num_filters=num_filters, reduction_ratio=16, name=name + '_fc')
def forward(self, x):
residual = x
......@@ -428,17 +353,11 @@ class SELayer(nn.Layer):
med_ch = int(num_channels / reduction_ratio)
stdv = 1.0 / math.sqrt(num_channels * 1.0)
self.squeeze = nn.Linear(
num_channels,
med_ch,
weight_attr=paddle.ParamAttr(
initializer=nn.initializer.Uniform(-stdv, stdv)))
num_channels, med_ch, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Uniform(-stdv, stdv)))
stdv = 1.0 / math.sqrt(med_ch * 1.0)
self.excitation = nn.Linear(
med_ch,
num_filters,
weight_attr=paddle.ParamAttr(
initializer=nn.initializer.Uniform(-stdv, stdv)))
med_ch, num_filters, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Uniform(-stdv, stdv)))
def forward(self, x):
pool = self.pool2d_gap(x)
......@@ -447,8 +366,7 @@ class SELayer(nn.Layer):
squeeze = F.relu(squeeze)
excitation = self.excitation(squeeze)
excitation = F.sigmoid(excitation)
excitation = paddle.reshape(
excitation, shape=[-1, self._num_channels, 1, 1])
excitation = paddle.reshape(excitation, shape=[-1, self._num_channels, 1, 1])
out = x * excitation
return out
......@@ -512,11 +430,7 @@ class HighResolutionModule(nn.Layer):
super(HighResolutionModule, self).__init__()
self.branches_func = Branches(
num_blocks=num_blocks,
in_channels=num_channels,
out_channels=num_filters,
has_se=has_se,
name=name)
num_blocks=num_blocks, in_channels=num_channels, out_channels=num_filters, has_se=has_se, name=name)
self.fuse_func = FuseLayers(
in_channels=num_filters,
......@@ -532,12 +446,7 @@ class HighResolutionModule(nn.Layer):
class FuseLayers(nn.Layer):
def __init__(self,
in_channels,
out_channels,
multi_scale_output=True,
name=None,
align_corners=False):
def __init__(self, in_channels, out_channels, multi_scale_output=True, name=None, align_corners=False):
super(FuseLayers, self).__init__()
self._actual_ch = len(in_channels) if multi_scale_output else 1
......@@ -562,8 +471,7 @@ class FuseLayers(nn.Layer):
for k in range(i - j):
if k == i - j - 1:
residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}_{}".format(
name, i + 1, j + 1, k + 1),
"residual_{}_layer_{}_{}_{}".format(name, i + 1, j + 1, k + 1),
ConvBN(
in_channels=pre_num_filters,
out_channels=out_channels[i],
......@@ -574,8 +482,7 @@ class FuseLayers(nn.Layer):
pre_num_filters = out_channels[i]
else:
residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}_{}".format(
name, i + 1, j + 1, k + 1),
"residual_{}_layer_{}_{}_{}".format(name, i + 1, j + 1, k + 1),
ConvBNReLU(
in_channels=pre_num_filters,
out_channels=out_channels[j],
......@@ -597,11 +504,7 @@ class FuseLayers(nn.Layer):
y = self.residual_func_list[residual_func_idx](x[j])
residual_func_idx += 1
y = F.interpolate(
y,
residual_shape,
mode='bilinear',
align_corners=self.align_corners)
y = F.interpolate(y, residual_shape, mode='bilinear', align_corners=self.align_corners)
residual = residual + y
elif j < i:
y = x[j]
......
......@@ -26,16 +26,10 @@ def SyncBatchNorm(*args, **kwargs):
class ConvBNReLU(nn.Layer):
def __init__(self,
in_channels,
out_channels,
kernel_size,
padding='same',
**kwargs):
def __init__(self, in_channels, out_channels, kernel_size, padding='same', **kwargs):
super().__init__()
self._conv = nn.Conv2D(
in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._batch_norm = SyncBatchNorm(out_channels)
......@@ -47,19 +41,12 @@ class ConvBNReLU(nn.Layer):
class ConvBN(nn.Layer):
def __init__(self,
in_channels,
out_channels,
kernel_size,
padding='same',
**kwargs):
def __init__(self, in_channels, out_channels, kernel_size, padding='same', **kwargs):
super().__init__()
self._conv = nn.Conv2D(
in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._batch_norm = SyncBatchNorm(out_channels)
def forward(self, x):
x = self._conv(x)
x = self._batch_norm(x)
return x
......@@ -6,13 +6,14 @@ import numpy as np
from FCN_HRNet_W18_Face_Seg.model import FCN, HRNet_W18
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="FCN_HRNet_W18_Face_Seg", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="FCN_HRNet_W18_Face_Seg", # 模型介绍
version="1.0.0" # 版本号
name="FCN_HRNet_W18_Face_Seg", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="FCN_HRNet_W18_Face_Seg", # 模型介绍
version="1.0.0" # 版本号
)
class FCN_HRNet_W18_Face_Seg(nn.Layer):
def __init__(self):
......@@ -41,10 +42,10 @@ class FCN_HRNet_W18_Face_Seg(nn.Layer):
if images is not None:
datas = images
# 返回数据列表
return datas
# 数据预处理函数
@staticmethod
def preprocess(images, batch_size):
......@@ -64,19 +65,18 @@ class FCN_HRNet_W18_Face_Seg(nn.Layer):
# 数据切分
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
return input_datas
# 结果归一化函数
@staticmethod
def normPRED(d):
ma = np.max(d)
mi = np.min(d)
dn = (d-mi)/(ma-mi)
dn = (d - mi) / (ma - mi)
return dn
......@@ -86,7 +86,7 @@ class FCN_HRNet_W18_Face_Seg(nn.Layer):
if visualization:
if not os.path.exists(output_dir):
os.mkdir(output_dir)
results = []
for output, image, i in zip(outputs, datas, range(len(datas))):
......@@ -96,24 +96,21 @@ class FCN_HRNet_W18_Face_Seg(nn.Layer):
# 图像缩放
h, w = image.shape[:2]
mask = cv2.resize(pred, (w, h))
mask[mask<0.5] = 0.
mask[mask>0.55] = 1.
mask[mask < 0.5] = 0.
mask[mask > 0.55] = 1.
# 计算输出图像
face = (image * mask[..., np.newaxis] + (1 - mask[..., np.newaxis]) * 255).astype(np.uint8)
# 格式还原
mask = (mask * 255).astype(np.uint8)
# 可视化结果保存
if visualization:
cv2.imwrite(os.path.join(output_dir, 'result_mask_%d.png' % i), mask)
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), face)
results.append({
'mask': mask,
'face': face
})
results.append({'mask': mask, 'face': face})
return results
......@@ -132,25 +129,19 @@ class FCN_HRNet_W18_Face_Seg(nn.Layer):
outputs = np.concatenate(outputs, 0)
return outputs
return outputs
def Segmentation(
self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
def Segmentation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
# 获取输入数据
datas = self.load_datas(paths, images)
# 数据预处理
input_datas = self.preprocess(datas, batch_size)
# 模型预测
outputs = self.predict(input_datas)
# 结果后处理
results = self.postprocess(outputs, datas, output_dir, visualization)
return results
\ No newline at end of file
return results
......@@ -6,18 +6,19 @@ from U2Net.u2net import U2NET
from U2Net.processor import Processor
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="U2Net", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Net", # 模型介绍
version="1.0.0" # 版本号
name="U2Net", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Net", # 模型介绍
version="1.0.0" # 版本号
)
class U2Net(nn.Layer):
def __init__(self):
super(U2Net, self).__init__()
self.model = U2NET(3,1)
self.model = U2NET(3, 1)
state_dict = paddle.load(os.path.join(self.directory, 'u2net.pdparams'))
self.model.set_dict(state_dict)
self.model.eval()
......@@ -26,21 +27,20 @@ class U2Net(nn.Layer):
outputs = []
for data in input_datas:
data = paddle.to_tensor(data, 'float32')
d1,d2,d3,d4,d5,d6,d7= self.model(data)
d1, d2, d3, d4, d5, d6, d7 = self.model(data)
outputs.append(d1.numpy())
outputs = np.concatenate(outputs, 0)
return outputs
def Segmentation(
self,
images=None,
paths=None,
batch_size=1,
input_size=320,
output_dir='output',
visualization=False):
def Segmentation(self,
images=None,
paths=None,
batch_size=1,
input_size=320,
output_dir='output',
visualization=False):
# 初始化数据处理器
processor = Processor(paths, images, batch_size, input_size)
......
......@@ -4,6 +4,7 @@ import numpy as np
__all__ = ['Processor']
class Processor():
def __init__(self, paths, images, batch_size, input_size):
# 图像列表
......@@ -25,7 +26,7 @@ class Processor():
if images is not None:
datas = images
# 返回数据列表
return datas
......@@ -34,22 +35,22 @@ class Processor():
input_datas = []
for image in imgs:
image = cv2.resize(image, (input_size, input_size))
tmpImg = np.zeros((image.shape[0],image.shape[1],3))
image = image/np.max(image)
tmpImg = np.zeros((image.shape[0], image.shape[1], 3))
image = image / np.max(image)
tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225
tmpImg[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
tmpImg[:, :, 1] = (image[:, :, 1] - 0.456) / 0.224
tmpImg[:, :, 2] = (image[:, :, 2] - 0.406) / 0.225
# convert BGR to RGB
tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = tmpImg[np.newaxis,:,:,:]
tmpImg = tmpImg[np.newaxis, :, :, :]
input_datas.append(tmpImg)
input_datas = np.concatenate(input_datas, 0)
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
......@@ -59,7 +60,7 @@ class Processor():
ma = np.max(d)
mi = np.min(d)
dn = (d-mi)/(ma-mi)
dn = (d - mi) / (ma - mi)
return dn
......@@ -68,10 +69,10 @@ class Processor():
results = []
if visualization and not os.path.exists(output_dir):
os.mkdir(output_dir)
for i, image in enumerate(self.imgs):
# normalization
pred = outputs[i,0,:,:]
pred = outputs[i, 0, :, :]
pred = self.normPRED(pred)
......@@ -79,17 +80,14 @@ class Processor():
h, w = image.shape[:2]
mask = cv2.resize(pred, (w, h))
output_img = (image*mask[..., np.newaxis] + (1-mask[..., np.newaxis])*255).astype(np.uint8)
output_img = (image * mask[..., np.newaxis] + (1 - mask[..., np.newaxis]) * 255).astype(np.uint8)
mask = (mask * 255).astype(np.uint8)
mask = (mask*255).astype(np.uint8)
if visualization:
cv2.imwrite(os.path.join(output_dir, 'result_mask_%d.png' % i), mask)
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), output_img)
results.append({
'mask': mask,
'front': output_img
})
return results
\ No newline at end of file
results.append({'mask': mask, 'front': output_img})
return results
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self, in_ch=3, out_ch=3, dirate=1):
super(REBNCONV, self).__init__()
self.conv_s1 = nn.Conv2D(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self, x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src, tar):
src = F.upsample(src, size=tar.shape[2:], mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer): #UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool5 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7, hx6), 1))
hx6dup = _upsample_like(hx6d, hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer): #UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer): #UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer): #UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer): #UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx2d = self.rebnconv2d(paddle.concat((hx3d, hx2), 1))
hx1d = self.rebnconv1d(paddle.concat((hx2d, hx1), 1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NET, self).__init__()
self.stage1 = RSU7(in_ch, 32, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 32, 128)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(128, 64, 256)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(256, 128, 512)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(512, 256, 512)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(512, 256, 512)
# decoder
self.stage5d = RSU4F(1024, 256, 512)
self.stage4d = RSU4(1024, 128, 256)
self.stage3d = RSU5(512, 64, 128)
self.stage2d = RSU6(256, 32, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(128, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(256, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(512, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(512, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NETP, self).__init__()
self.stage1 = RSU7(in_ch, 16, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 16, 64)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(64, 16, 64)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(64, 16, 64)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(64, 16, 64)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(64, 16, 64)
# decoder
self.stage5d = RSU4F(128, 16, 64)
self.stage4d = RSU4(128, 16, 64)
self.stage3d = RSU5(128, 16, 64)
self.stage2d = RSU6(128, 16, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(64, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
......@@ -6,18 +6,19 @@ from U2Netp.u2net import U2NETP
from U2Netp.processor import Processor
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="U2Netp", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Netp", # 模型介绍
version="1.0.0" # 版本号
name="U2Netp", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Netp", # 模型介绍
version="1.0.0" # 版本号
)
class U2Netp(nn.Layer):
def __init__(self):
super(U2Netp, self).__init__()
self.model = U2NETP(3,1)
self.model = U2NETP(3, 1)
state_dict = paddle.load(os.path.join(self.directory, 'u2netp.pdparams'))
self.model.set_dict(state_dict)
self.model.eval()
......@@ -26,21 +27,20 @@ class U2Netp(nn.Layer):
outputs = []
for data in input_datas:
data = paddle.to_tensor(data, 'float32')
d1,d2,d3,d4,d5,d6,d7= self.model(data)
d1, d2, d3, d4, d5, d6, d7 = self.model(data)
outputs.append(d1.numpy())
outputs = np.concatenate(outputs, 0)
return outputs
def Segmentation(
self,
images=None,
paths=None,
batch_size=1,
input_size=320,
output_dir='output',
visualization=False):
def Segmentation(self,
images=None,
paths=None,
batch_size=1,
input_size=320,
output_dir='output',
visualization=False):
# 初始化数据处理器
processor = Processor(paths, images, batch_size, input_size)
......
......@@ -4,6 +4,7 @@ import numpy as np
__all__ = ['Processor']
class Processor():
def __init__(self, paths, images, batch_size, input_size):
# 图像列表
......@@ -25,7 +26,7 @@ class Processor():
if images is not None:
datas = images
# 返回数据列表
return datas
......@@ -34,22 +35,22 @@ class Processor():
input_datas = []
for image in imgs:
image = cv2.resize(image, (input_size, input_size))
tmpImg = np.zeros((image.shape[0],image.shape[1],3))
image = image/np.max(image)
tmpImg = np.zeros((image.shape[0], image.shape[1], 3))
image = image / np.max(image)
tmpImg[:,:,0] = (image[:,:,0]-0.485)/0.229
tmpImg[:,:,1] = (image[:,:,1]-0.456)/0.224
tmpImg[:,:,2] = (image[:,:,2]-0.406)/0.225
tmpImg[:, :, 0] = (image[:, :, 0] - 0.485) / 0.229
tmpImg[:, :, 1] = (image[:, :, 1] - 0.456) / 0.224
tmpImg[:, :, 2] = (image[:, :, 2] - 0.406) / 0.225
# convert BGR to RGB
tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = tmpImg[np.newaxis,:,:,:]
tmpImg = tmpImg[np.newaxis, :, :, :]
input_datas.append(tmpImg)
input_datas = np.concatenate(input_datas, 0)
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
......@@ -59,7 +60,7 @@ class Processor():
ma = np.max(d)
mi = np.min(d)
dn = (d-mi)/(ma-mi)
dn = (d - mi) / (ma - mi)
return dn
......@@ -68,10 +69,10 @@ class Processor():
results = []
if visualization and not os.path.exists(output_dir):
os.mkdir(output_dir)
for i, image in enumerate(self.imgs):
# normalization
pred = outputs[i,0,:,:]
pred = outputs[i, 0, :, :]
pred = self.normPRED(pred)
......@@ -79,17 +80,14 @@ class Processor():
h, w = image.shape[:2]
mask = cv2.resize(pred, (w, h))
output_img = (image*mask[..., np.newaxis] + (1-mask[..., np.newaxis])*255).astype(np.uint8)
output_img = (image * mask[..., np.newaxis] + (1 - mask[..., np.newaxis]) * 255).astype(np.uint8)
mask = (mask * 255).astype(np.uint8)
mask = (mask*255).astype(np.uint8)
if visualization:
cv2.imwrite(os.path.join(output_dir, 'result_mask_%d.png' % i), mask)
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), output_img)
results.append({
'mask': mask,
'front': output_img
})
return results
\ No newline at end of file
results.append({'mask': mask, 'front': output_img})
return results
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self, in_ch=3, out_ch=3, dirate=1):
super(REBNCONV, self).__init__()
self.conv_s1 = nn.Conv2D(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self, x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src, tar):
src = F.upsample(src, size=tar.shape[2:], mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer): #UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool5 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7, hx6), 1))
hx6dup = _upsample_like(hx6d, hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer): #UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer): #UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer): #UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer): #UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx2d = self.rebnconv2d(paddle.concat((hx3d, hx2), 1))
hx1d = self.rebnconv1d(paddle.concat((hx2d, hx1), 1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NET, self).__init__()
self.stage1 = RSU7(in_ch, 32, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 32, 128)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(128, 64, 256)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(256, 128, 512)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(512, 256, 512)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(512, 256, 512)
# decoder
self.stage5d = RSU4F(1024, 256, 512)
self.stage4d = RSU4(1024, 128, 256)
self.stage3d = RSU5(512, 64, 128)
self.stage2d = RSU6(256, 32, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(128, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(256, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(512, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(512, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NETP, self).__init__()
self.stage1 = RSU7(in_ch, 16, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 16, 64)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(64, 16, 64)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(64, 16, 64)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(64, 16, 64)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(64, 16, 64)
# decoder
self.stage5d = RSU4F(128, 16, 64)
self.stage4d = RSU4(128, 16, 64)
self.stage3d = RSU5(128, 16, 64)
self.stage2d = RSU6(128, 16, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(64, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
......@@ -2,11 +2,12 @@ import math
import paddle
import paddle.nn as nn
class MLP(nn.Layer):
def __init__(self, embedding_size):
super(MLP, self).__init__()
self.dense_h_to_4h = nn.Linear(embedding_size, embedding_size*4)
self.dense_4h_to_h = nn.Linear(embedding_size*4, embedding_size)
self.dense_h_to_4h = nn.Linear(embedding_size, embedding_size * 4)
self.dense_4h_to_h = nn.Linear(embedding_size * 4, embedding_size)
self.act = nn.functional.gelu
def forward(self, x):
......@@ -14,12 +15,9 @@ class MLP(nn.Layer):
h2 = self.dense_4h_to_h(h)
return h2
class Attention(nn.Layer):
def __init__(self,
embedding_size,
num_attention_heads,
attention_dropout,
residual_dropout):
def __init__(self, embedding_size, num_attention_heads, attention_dropout, residual_dropout):
super(Attention, self).__init__()
self.num_attention_heads = num_attention_heads
self.size_per_head = embedding_size // num_attention_heads
......@@ -38,11 +36,11 @@ class Attention(nn.Layer):
self.seq_len = x.shape[1]
x = self.query_key_value(x)
q, k, v = x.split(num_or_sections=3, axis=2)
q = self.split_heads(q)
k = self.split_heads(k)
v = self.split_heads(v)
if kv_cache is not None:
pk, pv = paddle.unstack(kv_cache, axis=1)
k = paddle.concat([pk, k], axis=-2)
......@@ -69,12 +67,9 @@ class Attention(nn.Layer):
return y, cached_kv
class Block(nn.Layer):
def __init__(self,
embedding_size,
num_attention_heads,
attention_dropout,
residual_dropout):
def __init__(self, embedding_size, num_attention_heads, attention_dropout, residual_dropout):
super(Block, self).__init__()
self.input_layernorm = nn.LayerNorm(embedding_size, epsilon=1e-5)
self.attention = Attention(embedding_size, num_attention_heads, attention_dropout, residual_dropout)
......@@ -89,24 +84,17 @@ class Block(nn.Layer):
x = x + z
return x, cached_kv
class Transformer(nn.Layer):
def __init__(self,
layer_size,
embedding_size,
num_attention_heads,
attention_dropout,
residual_dropout):
def __init__(self, layer_size, embedding_size, num_attention_heads, attention_dropout, residual_dropout):
super(Transformer, self).__init__()
self.layers = nn.LayerList([Block(
embedding_size,
num_attention_heads,
attention_dropout,
residual_dropout)
for _ in range(layer_size)])
self.layers = nn.LayerList([
Block(embedding_size, num_attention_heads, attention_dropout, residual_dropout) for _ in range(layer_size)
])
self.final_layernorm = nn.LayerNorm(embedding_size, epsilon=1e-5)
def forward(self, x, kv_cache=None):
cached_kvs = []
for i, layer in enumerate(self.layers):
......@@ -115,26 +103,16 @@ class Transformer(nn.Layer):
x = self.final_layernorm(x)
return x, paddle.stack(cached_kvs)
class GPT2Model(nn.Layer):
def __init__(self,
vocab_size,
layer_size,
block_size,
embedding_dropout,
embedding_size,
num_attention_heads,
attention_dropout,
residual_dropout):
def __init__(self, vocab_size, layer_size, block_size, embedding_dropout, embedding_size, num_attention_heads,
attention_dropout, residual_dropout):
super(GPT2Model, self).__init__()
self.word_embeddings = nn.Embedding(vocab_size, embedding_size)
self.position_embeddings = nn.Embedding(block_size, embedding_size)
self.emb_drop = nn.Dropout(embedding_dropout)
self.transformer = Transformer(
layer_size,
embedding_size,
num_attention_heads,
attention_dropout,
residual_dropout)
self.transformer = Transformer(layer_size, embedding_size, num_attention_heads, attention_dropout,
residual_dropout)
def forward(self, x, kv_cache=None, use_cache=False):
if kv_cache is None:
......@@ -155,16 +133,17 @@ class GPT2Model(nn.Layer):
else:
return x
if __name__ == '__main__':
gpt = GPT2Model(
vocab_size=30000,
layer_size=2,
block_size=1024,
embedding_dropout=0.0,
embedding_size=2560,
num_attention_heads=32,
attention_dropout=0.0,
residual_dropout=0.0)
vocab_size=30000,
layer_size=2,
block_size=1024,
embedding_dropout=0.0,
embedding_size=2560,
num_attention_heads=32,
attention_dropout=0.0,
residual_dropout=0.0)
gpt.eval()
out, cached_kvs = gpt(paddle.ones([1,1], 'int64'), paddle.randn([32, 1, 2, 32, 9, 80], 'float32'), use_cache=True)
out, cached_kvs = gpt(paddle.ones([1, 1], 'int64'), paddle.randn([32, 1, 2, 32, 9, 80], 'float32'), use_cache=True)
print(out.shape, cached_kvs.shape)
......@@ -13,8 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tokenization classes for OpenAI GPT."""
from __future__ import (absolute_import, division, print_function,
unicode_literals)
from __future__ import (absolute_import, division, print_function, unicode_literals)
import json
from io import open
......@@ -29,11 +28,12 @@ except ImportError:
def lru_cache():
return lambda func: func
class GPT2Tokenizer(object):
def __init__(self, vocab_file, model_file, max_len=None):
self.max_len = max_len if max_len is not None else int(1e12)
self.encoder = json.load(open(vocab_file))
self.decoder = {v:k for k,v in self.encoder.items()}
self.decoder = {v: k for k, v in self.encoder.items()}
self.sp = spm.SentencePieceProcessor(model_file=model_file)
self.translator = str.maketrans(" \n", "\u2582\u2583")
......
......@@ -8,8 +8,8 @@ CPM-LM 是一个基于 GPT-2 的预训练生成模型,参数规模达 26 亿
## API
```python
def greedy_search(
text,
max_len=32,
text,
max_len=32,
end_word=None):
```
文本生成 API ,根据输入的文字进行文本生成,使用 Greedy Search 进行解码,生成的文本单一且确定,适合于问答类的文本生成。
......@@ -24,12 +24,12 @@ def greedy_search(
```python
def sample(
text,
max_len=32,
end_word=None,
repitition_penalty=1.0,
temperature=1.0,
top_k=0,
text,
max_len=32,
end_word=None,
repitition_penalty=1.0,
temperature=1.0,
top_k=0,
top_p=1.0):
```
文本生成 API ,根据输入的文字进行文本生成,使用采样的方式进行解码,生成的文本比较多样,适合于文章类的文本生成。
......@@ -137,4 +137,4 @@ https://github.com/jm12138/CPM-Generate-Paddle
## 依赖
paddlepaddle >= 2.0.0rc0
paddlehub >= 2.0.0b1
\ No newline at end of file
paddlehub >= 2.0.0b1
......@@ -5,13 +5,14 @@ from paddle.nn import Layer
from paddlehub.module.module import moduleinfo
from CPM_LM.GPT2 import GPT2Model, GPT2Tokenizer
@moduleinfo(
name="CPM_LM", # 模型名称
type="NLP/NLG", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="CPM_LM", # 模型介绍
version="1.0.1" # 版本号
name="CPM_LM", # 模型名称
type="NLP/NLG", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="CPM_LM", # 模型介绍
version="1.0.1" # 版本号
)
class CPM_LM(Layer):
def __init__(self, max_len=512):
......@@ -56,7 +57,7 @@ class CPM_LM(Layer):
if end_word is not None:
stop_id = self.tokenizer.encode(end_word)
length = len(stop_id)
# 初始预测
ids = self.tokenizer.encode(text)
input_id = paddle.to_tensor(np.array(ids).reshape(1, -1).astype('int64'))
......@@ -65,17 +66,17 @@ class CPM_LM(Layer):
ids.append(next_token)
# 使用缓存进行继续预测
for i in range(max_len-1):
for i in range(max_len - 1):
input_id = paddle.to_tensor(np.array([next_token]).reshape(1, -1).astype('int64'))
output, cached_kvs = self.model(input_id, cached_kvs, use_cache=True)
next_token = int(np.argmax(output[0, -1].numpy()))
ids.append(next_token)
if next_token==end_id:
if next_token == end_id:
break
# 根据终止标志停止预测
if (end_word is not None) and (ids[-length:]==stop_id):
if (end_word is not None) and (ids[-length:] == stop_id):
break
return self.tokenizer.decode(ids)
......@@ -135,7 +136,7 @@ class CPM_LM(Layer):
ids += [int(next_token)]
# 使用缓存进行继续预测
for i in range(max_len-1):
for i in range(max_len - 1):
input_id = paddle.to_tensor(np.array([next_token]).reshape(1, -1).astype('int64'))
output, cached_kvs = self.model(input_id, cached_kvs, use_cache=True)
next_token_logits = output[0, -1, :]
......@@ -144,14 +145,15 @@ class CPM_LM(Layer):
next_token_logits = next_token_logits / temperature
next_token_logits[self.tokenizer.encoder['<unk>']] = -float('Inf')
filtered_logits = self.top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
next_token = paddle.multinomial(paddle.nn.functional.softmax(filtered_logits, axis=-1), num_samples=1).numpy()
next_token = paddle.multinomial(
paddle.nn.functional.softmax(filtered_logits, axis=-1), num_samples=1).numpy()
ids += [int(next_token)]
if next_token==end_id:
if next_token == end_id:
break
# 根据终止标志停止预测
if (end_word is not None) and (ids[-length:]==stop_id):
if (end_word is not None) and (ids[-length:] == stop_id):
break
return self.tokenizer.decode(ids)
\ No newline at end of file
return self.tokenizer.decode(ids)
......@@ -17,7 +17,7 @@ import sys
from easydict import EasyDict
__version__ = '2.0.0-beta0'
__version__ = '2.0.0-beta2'
from paddlehub import env
from paddlehub.config import config
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册