“26eb54eb37e0515f863243c133fe0a72bfd5c6af”上不存在“examples/transv1.8to2.x/deepspeech/utils/error_rate.py”
提交 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,6 +7,7 @@ import paddle.nn as nn ...@@ -7,6 +7,7 @@ import paddle.nn as nn
import paddlehub as hub import paddlehub as hub
from paddlehub.module.module import moduleinfo from paddlehub.module.module import moduleinfo
@moduleinfo( @moduleinfo(
name="ID_Photo_GEN", # 模型名称 name="ID_Photo_GEN", # 模型名称
type="CV", # 模型类型 type="CV", # 模型类型
...@@ -45,10 +46,7 @@ class ID_Photo_GEN(nn.Layer): ...@@ -45,10 +46,7 @@ class ID_Photo_GEN(nn.Layer):
# 数据预处理函数 # 数据预处理函数
def preprocess(self, images, batch_size, use_gpu): def preprocess(self, images, batch_size, use_gpu):
# 获取人脸关键点 # 获取人脸关键点
outputs = self.face_detector.keypoint_detection( outputs = self.face_detector.keypoint_detection(images=images, batch_size=batch_size, use_gpu=use_gpu)
images=images,
batch_size=batch_size,
use_gpu=use_gpu)
crops = [] crops = []
for output, image in zip(outputs, images): for output, image in zip(outputs, images):
...@@ -58,7 +56,8 @@ class ID_Photo_GEN(nn.Layer): ...@@ -58,7 +56,8 @@ class ID_Photo_GEN(nn.Layer):
# rotation angle # rotation angle
left_eye_corner = landmarks[36] left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45] 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 # image size after rotating
height, width, _ = image.shape height, width, _ = image.shape
...@@ -103,14 +102,15 @@ class ID_Photo_GEN(nn.Layer): ...@@ -103,14 +102,15 @@ class ID_Photo_GEN(nn.Layer):
h, w = image.shape[:2] h, w = image.shape[:2]
left_white = max(0, -left) left_white = max(0, -left)
left = max(0, left) left = max(0, left)
right = min(right, w-1) right = min(right, w - 1)
right_white = left_white + (right-left) right_white = left_white + (right - left)
top_white = max(0, -top) top_white = max(0, -top)
top = max(0, top) top = max(0, top)
bottom = min(bottom, h-1) bottom = min(bottom, h - 1)
bottom_white = top_white + (bottom - top) 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) crops.append(image_crop)
# 获取人像分割的输出 # 获取人像分割的输出
...@@ -168,22 +168,11 @@ class ID_Photo_GEN(nn.Layer): ...@@ -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, 'blue_%d.jpg' % i), blue)
cv2.imwrite(os.path.join(output_dir, 'red_%d.jpg' % i), red) cv2.imwrite(os.path.join(output_dir, 'red_%d.jpg' % i), red)
results.append({ results.append({'write': write, 'blue': blue, 'red': red})
'write': write,
'blue': blue,
'red': red
})
return results return results
def Photo_GEN( def Photo_GEN(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False, use_gpu=False):
self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
# 获取输入数据 # 获取输入数据
images = self.load_datas(paths, images) images = self.load_datas(paths, images)
......
...@@ -17,72 +17,66 @@ class ResnetGenerator(nn.Layer): ...@@ -17,72 +17,66 @@ class ResnetGenerator(nn.Layer):
nn.ReLU() nn.ReLU()
] ]
DownBlock += [ DownBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf)]
HourGlass(ngf, ngf),
HourGlass(ngf, ngf)
]
# Down-Sampling # Down-Sampling
n_downsampling = 2 n_downsampling = 2
for i in range(n_downsampling): for i in range(n_downsampling):
mult = 2 ** i mult = 2**i
DownBlock += [ DownBlock += [
nn.Pad2D([1, 1, 1, 1], 'reflect'), nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf*mult, ngf*mult*2, kernel_size=3, stride=2, 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.InstanceNorm2D(ngf * mult * 2, weight_attr=False, bias_attr=False),
nn.ReLU() nn.ReLU()
] ]
# Encoder Bottleneck # Encoder Bottleneck
mult = 2 ** n_downsampling mult = 2**n_downsampling
for i in range(n_blocks): 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 # Class Activation Map
self.gap_fc = nn.Linear(ngf*mult, 1, bias_attr=False) self.gap_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.gmp_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.conv1x1 = nn.Conv2D(ngf * mult * 2, ngf * mult, kernel_size=1, stride=1)
self.relu = nn.ReLU() self.relu = nn.ReLU()
# Gamma, Beta block # Gamma, Beta block
FC = [] FC = []
if self.light: if self.light:
FC += [ FC += [
nn.Linear(ngf*mult, ngf*mult, bias_attr=False), nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(), nn.ReLU(),
nn.Linear(ngf*mult, ngf*mult, bias_attr=False), nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU() nn.ReLU()
] ]
else: else:
FC += [ 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.ReLU(),
nn.Linear(ngf*mult, ngf*mult, bias_attr=False), nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU() nn.ReLU()
] ]
# Decoder Bottleneck # Decoder Bottleneck
mult = 2 ** n_downsampling mult = 2**n_downsampling
for i in range(n_blocks): 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 # Up-Sampling
UpBlock = [] UpBlock = []
for i in range(n_downsampling): for i in range(n_downsampling):
mult = 2 ** (n_downsampling - i) mult = 2**(n_downsampling - i)
UpBlock += [ UpBlock += [
nn.Upsample(scale_factor=2), nn.Upsample(scale_factor=2),
nn.Pad2D([1, 1, 1, 1], 'reflect'), nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf*mult, ngf*mult//2, kernel_size=3, stride=1, bias_attr=False), nn.Conv2D(ngf * mult, ngf * mult // 2, kernel_size=3, stride=1, bias_attr=False),
LIN(ngf*mult//2), LIN(ngf * mult // 2),
nn.ReLU() nn.ReLU()
] ]
UpBlock += [ UpBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf, False)]
HourGlass(ngf, ngf),
HourGlass(ngf, ngf, False)
]
UpBlock += [ UpBlock += [
nn.Pad2D([3, 3, 3, 3], 'reflect'), nn.Pad2D([3, 3, 3, 3], 'reflect'),
...@@ -101,7 +95,7 @@ class ResnetGenerator(nn.Layer): ...@@ -101,7 +95,7 @@ class ResnetGenerator(nn.Layer):
content_features = [] content_features = []
for i in range(self.n_blocks): 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])) content_features.append(F.adaptive_avg_pool2d(x, 1).reshape([bs, -1]))
gap = F.adaptive_avg_pool2d(x, 1) gap = F.adaptive_avg_pool2d(x, 1)
...@@ -127,7 +121,7 @@ class ResnetGenerator(nn.Layer): ...@@ -127,7 +121,7 @@ class ResnetGenerator(nn.Layer):
style_features = self.FC(x.reshape([bs, -1])) style_features = self.FC(x.reshape([bs, -1]))
for i in range(self.n_blocks): 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) out = self.UpBlock(x)
...@@ -140,25 +134,20 @@ class ConvBlock(nn.Layer): ...@@ -140,25 +134,20 @@ class ConvBlock(nn.Layer):
self.dim_in = dim_in self.dim_in = dim_in
self.dim_out = dim_out self.dim_out = dim_out
self.conv_block1 = self.__convblock(dim_in, dim_out//2) self.conv_block1 = self.__convblock(dim_in, dim_out // 2)
self.conv_block2 = self.__convblock(dim_out//2, dim_out//4) self.conv_block2 = self.__convblock(dim_out // 2, dim_out // 4)
self.conv_block3 = self.__convblock(dim_out//4, dim_out//4) self.conv_block3 = self.__convblock(dim_out // 4, dim_out // 4)
if self.dim_in != self.dim_out: if self.dim_in != self.dim_out:
self.conv_skip = nn.Sequential( self.conv_skip = nn.Sequential(
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.ReLU(),
nn.ReLU(), nn.Conv2D(dim_in, dim_out, kernel_size=1, stride=1, bias_attr=False))
nn.Conv2D(dim_in, dim_out, kernel_size=1, stride=1, bias_attr=False)
)
@staticmethod @staticmethod
def __convblock(dim_in, dim_out): def __convblock(dim_in, dim_out):
return nn.Sequential( return nn.Sequential(
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.ReLU(), nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.ReLU(), nn.Conv2D(dim_in, dim_out, kernel_size=3, stride=1, bias_attr=False))
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): def forward(self, x):
residual = x residual = x
...@@ -182,24 +171,24 @@ class HourGlassBlock(nn.Layer): ...@@ -182,24 +171,24 @@ class HourGlassBlock(nn.Layer):
self.n_block = 9 self.n_block = 9
for i in range(self.n_skip): 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): 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): def forward(self, x):
skips = [] skips = []
for i in range(self.n_skip): 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 = 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) x = self.ConvBlock5(x)
for i in range(self.n_skip): 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 = F.upsample(x, scale_factor=2)
x = skips[self.n_skip-i-1] + x x = skips[self.n_skip - i - 1] + x
return x return x
...@@ -210,12 +199,9 @@ class HourGlass(nn.Layer): ...@@ -210,12 +199,9 @@ class HourGlass(nn.Layer):
self.use_res = use_res self.use_res = use_res
self.HG = nn.Sequential( self.HG = nn.Sequential(
HourGlassBlock(dim_in), HourGlassBlock(dim_in), ConvBlock(dim_out, dim_out),
ConvBlock(dim_out, dim_out),
nn.Conv2D(dim_out, dim_out, kernel_size=1, stride=1, bias_attr=False), 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.InstanceNorm2D(dim_out, weight_attr=False, bias_attr=False), nn.ReLU())
nn.ReLU()
)
self.Conv1 = nn.Conv2D(dim_out, 3, kernel_size=1, stride=1) self.Conv1 = nn.Conv2D(dim_out, 3, kernel_size=1, stride=1)
...@@ -292,11 +278,11 @@ class SoftAdaLIN(nn.Layer): ...@@ -292,11 +278,11 @@ class SoftAdaLIN(nn.Layer):
self.w_gamma = self.create_parameter([1, num_features], default_initializer=nn.initializer.Constant(0.)) 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.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), self.c_gamma = nn.Sequential(
nn.ReLU(), nn.Linear(num_features, num_features, bias_attr=False), nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False)) nn.Linear(num_features, num_features, bias_attr=False))
self.c_beta = nn.Sequential(nn.Linear(num_features, num_features, bias_attr=False), self.c_beta = nn.Sequential(
nn.ReLU(), nn.Linear(num_features, num_features, bias_attr=False), nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False)) nn.Linear(num_features, num_features, bias_attr=False))
self.s_gamma = 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) self.s_beta = nn.Linear(num_features, num_features, bias_attr=False)
......
...@@ -8,6 +8,7 @@ import paddlehub as hub ...@@ -8,6 +8,7 @@ import paddlehub as hub
from Photo2Cartoon.model import ResnetGenerator from Photo2Cartoon.model import ResnetGenerator
from paddlehub.module.module import moduleinfo from paddlehub.module.module import moduleinfo
@moduleinfo( @moduleinfo(
name="Photo2Cartoon", # 模型名称 name="Photo2Cartoon", # 模型名称
type="CV", # 模型类型 type="CV", # 模型类型
...@@ -56,10 +57,7 @@ class Photo2Cartoon(nn.Layer): ...@@ -56,10 +57,7 @@ class Photo2Cartoon(nn.Layer):
# 数据预处理函数 # 数据预处理函数
def preprocess(self, images, batch_size, use_gpu): def preprocess(self, images, batch_size, use_gpu):
# 获取人脸关键点 # 获取人脸关键点
outputs = self.face_detector.keypoint_detection( outputs = self.face_detector.keypoint_detection(images=images, batch_size=batch_size, use_gpu=use_gpu)
images=images,
batch_size=batch_size,
use_gpu=use_gpu)
crops = [] crops = []
for output, image in zip(outputs, images): for output, image in zip(outputs, images):
...@@ -69,7 +67,8 @@ class Photo2Cartoon(nn.Layer): ...@@ -69,7 +67,8 @@ class Photo2Cartoon(nn.Layer):
# rotation angle # rotation angle
left_eye_corner = landmarks[36] left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45] 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 # image size after rotating
height, width, _ = image.shape height, width, _ = image.shape
...@@ -114,14 +113,15 @@ class Photo2Cartoon(nn.Layer): ...@@ -114,14 +113,15 @@ class Photo2Cartoon(nn.Layer):
h, w = image.shape[:2] h, w = image.shape[:2]
left_white = max(0, -left) left_white = max(0, -left)
left = max(0, left) left = max(0, left)
right = min(right, w-1) right = min(right, w - 1)
right_white = left_white + (right-left) right_white = left_white + (right - left)
top_white = max(0, -top) top_white = max(0, -top)
top = max(0, top) top = max(0, top)
bottom = min(bottom, h-1) bottom = min(bottom, h - 1)
bottom_white = top_white + (bottom - top) 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) crops.append(image_crop)
# 获取人像分割的输出 # 获取人像分割的输出
...@@ -158,7 +158,7 @@ class Photo2Cartoon(nn.Layer): ...@@ -158,7 +158,7 @@ class Photo2Cartoon(nn.Layer):
# 切分数据 # 切分数据
datas_num = input_datas.shape[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) input_datas = np.array_split(input_datas, split_num)
return input_datas, masks return input_datas, masks
...@@ -207,8 +207,7 @@ class Photo2Cartoon(nn.Layer): ...@@ -207,8 +207,7 @@ class Photo2Cartoon(nn.Layer):
return cartoons return cartoons
def Cartoon_GEN( def Cartoon_GEN(self,
self,
images=None, images=None,
paths=None, paths=None,
batch_size=1, batch_size=1,
......
...@@ -6,6 +6,7 @@ from U2Net_Portrait.u2net import U2NET ...@@ -6,6 +6,7 @@ from U2Net_Portrait.u2net import U2NET
from U2Net_Portrait.processor import Processor from U2Net_Portrait.processor import Processor
from paddlehub.module.module import moduleinfo from paddlehub.module.module import moduleinfo
@moduleinfo( @moduleinfo(
name="U2Net_Portrait", # 模型名称 name="U2Net_Portrait", # 模型名称
type="CV", # 模型类型 type="CV", # 模型类型
...@@ -17,7 +18,7 @@ from paddlehub.module.module import moduleinfo ...@@ -17,7 +18,7 @@ from paddlehub.module.module import moduleinfo
class U2Net_Portrait(nn.Layer): class U2Net_Portrait(nn.Layer):
def __init__(self): def __init__(self):
super(U2Net_Portrait, self).__init__() 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')) state_dict = paddle.load(os.path.join(self.directory, 'u2net_portrait.pdparams'))
self.model.set_dict(state_dict) self.model.set_dict(state_dict)
self.model.eval() self.model.eval()
...@@ -26,15 +27,14 @@ class U2Net_Portrait(nn.Layer): ...@@ -26,15 +27,14 @@ class U2Net_Portrait(nn.Layer):
outputs = [] outputs = []
for data in input_datas: for data in input_datas:
data = paddle.to_tensor(data, 'float32') 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.append(d1.numpy())
outputs = np.concatenate(outputs, 0) outputs = np.concatenate(outputs, 0)
return outputs return outputs
def Portrait_GEN( def Portrait_GEN(self,
self,
images=None, images=None,
paths=None, paths=None,
scale=1, scale=1,
......
...@@ -5,6 +5,7 @@ import paddlehub as hub ...@@ -5,6 +5,7 @@ import paddlehub as hub
__all__ = ['Processor'] __all__ = ['Processor']
class Processor(): class Processor():
def __init__(self, paths, images, batch_size, face_detection=True, scale=1): def __init__(self, paths, images, batch_size, face_detection=True, scale=1):
# 图像列表 # 图像列表
...@@ -35,10 +36,7 @@ class Processor(): ...@@ -35,10 +36,7 @@ class Processor():
if face_detection: if face_detection:
# face detection # face detection
face_detector = hub.Module(name="pyramidbox_lite_mobile") face_detector = hub.Module(name="pyramidbox_lite_mobile")
results = face_detector.face_detection(images=imgs, results = face_detector.face_detection(images=imgs, use_gpu=False, visualization=False, confs_threshold=0.5)
use_gpu=False,
visualization=False,
confs_threshold=0.5)
im_faces = [] im_faces = []
for datas, img in zip(results, imgs): for datas, img in zip(results, imgs):
for face in datas['data']: for face in datas['data']:
...@@ -46,52 +44,64 @@ class Processor(): ...@@ -46,52 +44,64 @@ class Processor():
l, r, t, b = [face['left'], face['right'], face['top'], face['bottom']] l, r, t, b = [face['left'], face['right'], face['top'], face['bottom']]
# square crop # square crop
pad = max(int(scale*(r-l)), int(scale*(b-t))) pad = max(int(scale * (r - l)), int(scale * (b - t)))
c_w, c_h = (r-l)//2+l, (b-t)//2+t c_w, c_h = (r - l) // 2 + l, (b - t) // 2 + t
top = 0 if c_h-pad<0 else c_h-pad top = 0 if c_h - pad < 0 else c_h - pad
bottom = pad + c_h 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 right = pad + c_w
crop = img[top:bottom, left:right] crop = img[top:bottom, left:right]
# resize # 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) im_faces.append(im_face)
else: else:
im_faces = [] im_faces = []
for img in imgs: for img in imgs:
h, w = img.shape[:2] h, w = img.shape[:2]
if h>w: if h > w:
if (h-w)%2==0: 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))) 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: 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: else:
if (w-h)%2==0: 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))) 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: 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))) img = np.pad(
im_face = cv2.resize(img, (512,512), interpolation = cv2.INTER_AREA) 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) im_faces.append(im_face)
input_datas = [] input_datas = []
for im_face in im_faces: for im_face in im_faces:
tmpImg = np.zeros((im_face.shape[0],im_face.shape[1],3)) tmpImg = np.zeros((im_face.shape[0], im_face.shape[1], 3))
im_face = im_face/np.max(im_face) im_face = im_face / np.max(im_face)
tmpImg[:,:,0] = (im_face[:,:,2]-0.406)/0.225 tmpImg[:, :, 0] = (im_face[:, :, 2] - 0.406) / 0.225
tmpImg[:,:,1] = (im_face[:,:,1]-0.456)/0.224 tmpImg[:, :, 1] = (im_face[:, :, 1] - 0.456) / 0.224
tmpImg[:,:,2] = (im_face[:,:,0]-0.485)/0.229 tmpImg[:, :, 2] = (im_face[:, :, 0] - 0.485) / 0.229
# convert BGR to RGB # convert BGR to RGB
tmpImg = tmpImg.transpose((2, 0, 1)) tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = tmpImg[np.newaxis,:,:,:] tmpImg = tmpImg[np.newaxis, :, :, :]
input_datas.append(tmpImg) input_datas.append(tmpImg)
input_datas = np.concatenate(input_datas, 0) input_datas = np.concatenate(input_datas, 0)
datas_num = input_datas.shape[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) input_datas = np.array_split(input_datas, split_num)
...@@ -101,7 +111,7 @@ class Processor(): ...@@ -101,7 +111,7 @@ class Processor():
ma = np.max(d) ma = np.max(d)
mi = np.min(d) mi = np.min(d)
dn = (d-mi)/(ma-mi) dn = (d - mi) / (ma - mi)
return dn return dn
...@@ -113,13 +123,13 @@ class Processor(): ...@@ -113,13 +123,13 @@ class Processor():
for i in range(outputs.shape[0]): for i in range(outputs.shape[0]):
# normalization # normalization
pred = 1.0 - outputs[i,0,:,:] pred = 1.0 - outputs[i, 0, :, :]
pred = self.normPRED(pred) pred = self.normPRED(pred)
# convert torch tensor to numpy array # convert torch tensor to numpy array
pred = pred.squeeze() pred = pred.squeeze()
pred = (pred*255).astype(np.uint8) pred = (pred * 255).astype(np.uint8)
results.append(pred) results.append(pred)
......
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 ...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model'] __all__ = ['Model']
class Model(): class Model():
# 初始化函数 # 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True): def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
...@@ -24,7 +25,9 @@ class Model(): ...@@ -24,7 +25,9 @@ class Model():
try: try:
int(os.environ.get('CUDA_VISIBLE_DEVICES')) int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception: 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 use_gpu = False
# 加载模型参数 # 加载模型参数
......
...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config ...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model'] __all__ = ['Model']
class Model(): class Model():
# 初始化函数 # 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True): def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
...@@ -24,7 +25,9 @@ class Model(): ...@@ -24,7 +25,9 @@ class Model():
try: try:
int(os.environ.get('CUDA_VISIBLE_DEVICES')) int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception: 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 use_gpu = False
# 加载模型参数 # 加载模型参数
......
...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config ...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model'] __all__ = ['Model']
class Model(): class Model():
# 初始化函数 # 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True): def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
...@@ -24,7 +25,9 @@ class Model(): ...@@ -24,7 +25,9 @@ class Model():
try: try:
int(os.environ.get('CUDA_VISIBLE_DEVICES')) int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception: 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 use_gpu = False
# 加载模型参数 # 加载模型参数
......
...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config ...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model'] __all__ = ['Model']
class Model(): class Model():
# 初始化函数 # 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True): def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
...@@ -24,7 +25,9 @@ class Model(): ...@@ -24,7 +25,9 @@ class Model():
try: try:
int(os.environ.get('CUDA_VISIBLE_DEVICES')) int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception: 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 use_gpu = False
# 加载模型参数 # 加载模型参数
......
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 ...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['InferenceModel'] __all__ = ['InferenceModel']
class InferenceModel(): class InferenceModel():
# 初始化函数 # 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True): def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True):
...@@ -27,12 +28,8 @@ class InferenceModel(): ...@@ -27,12 +28,8 @@ class InferenceModel():
''' '''
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' % ( return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (len(
len(self.input_handles), self.input_handles), str(self.input_names), len(self.output_handles), str(self.output_names))
str(self.input_names),
len(self.output_handles),
str(self.output_names)
)
# 类调用函数 # 类调用函数
def __call__(self, *input_datas, batch_size=1): def __call__(self, *input_datas, batch_size=1):
...@@ -59,7 +56,9 @@ class InferenceModel(): ...@@ -59,7 +56,9 @@ class InferenceModel():
try: try:
int(os.environ.get('CUDA_VISIBLE_DEVICES')) int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception: 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 use_gpu = False
# 加载模型参数 # 加载模型参数
......
...@@ -26,12 +26,7 @@ class MiDaS_Large(Module): ...@@ -26,12 +26,7 @@ class MiDaS_Large(Module):
model_path = os.path.join(self.directory, "model-f6b98070") model_path = os.path.join(self.directory, "model-f6b98070")
# 加载模型 # 加载模型
self.model = InferenceModel( self.model = InferenceModel(modelpath=model_path, use_gpu=use_gpu, use_mkldnn=False, combined=True)
modelpath=model_path,
use_gpu=use_gpu,
use_mkldnn=False,
combined=True
)
self.model.eval() self.model.eval()
# 数据预处理配置 # 数据预处理配置
...@@ -46,8 +41,7 @@ class MiDaS_Large(Module): ...@@ -46,8 +41,7 @@ class MiDaS_Large(Module):
resize_method="upper_bound", resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC, image_interpolation_method=cv2.INTER_CUBIC,
), ),
NormalizeImage(mean=[0.485, 0.456, 0.406], NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
std=[0.229, 0.224, 0.225]),
PrepareForNet() PrepareForNet()
]) ])
...@@ -113,12 +107,7 @@ class MiDaS_Large(Module): ...@@ -113,12 +107,7 @@ class MiDaS_Large(Module):
return outputs return outputs
# 深度估计函数 # 深度估计函数
def depth_estimation(self, def depth_estimation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
# 加载数据 # 加载数据
datas = self.load_datas(paths, images) datas = self.load_datas(paths, images)
......
...@@ -7,6 +7,7 @@ import cv2 ...@@ -7,6 +7,7 @@ import cv2
class Resize(object): class Resize(object):
"""Resize sample to given size (width, height). """Resize sample to given size (width, height).
""" """
def __init__(self, def __init__(self,
width, width,
height, height,
...@@ -51,12 +52,10 @@ class Resize(object): ...@@ -51,12 +52,10 @@ class Resize(object):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val: if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) * y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
self.__multiple_of).astype(int)
if y < min_val: if y < min_val:
y = (np.ceil(x / self.__multiple_of) * y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
self.__multiple_of).astype(int)
return y return y
...@@ -91,31 +90,24 @@ class Resize(object): ...@@ -91,31 +90,24 @@ class Resize(object):
# fit height # fit height
scale_width = scale_height scale_width = scale_height
else: else:
raise ValueError( raise ValueError(f"resize_method {self.__resize_method} not implemented")
f"resize_method {self.__resize_method} not implemented")
if self.__resize_method == "lower_bound": if self.__resize_method == "lower_bound":
new_height = self.constrain_to_multiple_of(scale_height * height, new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
min_val=self.__height) new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
elif self.__resize_method == "upper_bound": elif self.__resize_method == "upper_bound":
new_height = self.constrain_to_multiple_of(scale_height * height, new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
max_val=self.__height) new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
elif self.__resize_method == "minimal": elif self.__resize_method == "minimal":
new_height = self.constrain_to_multiple_of(scale_height * height) new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width) new_width = self.constrain_to_multiple_of(scale_width * width)
else: else:
raise ValueError( raise ValueError(f"resize_method {self.__resize_method} not implemented")
f"resize_method {self.__resize_method} not implemented")
return (new_width, new_height) return (new_width, new_height)
def __call__(self, sample): def __call__(self, sample):
width, height = self.get_size(sample["image"].shape[1], width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
sample["image"].shape[0])
# resize sample # resize sample
sample["image"] = cv2.resize( sample["image"] = cv2.resize(
...@@ -133,8 +125,7 @@ class Resize(object): ...@@ -133,8 +125,7 @@ class Resize(object):
) )
if "depth" in sample: if "depth" in sample:
sample["depth"] = cv2.resize(sample["depth"], (width, height), sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
interpolation=cv2.INTER_NEAREST)
sample["mask"] = cv2.resize( sample["mask"] = cv2.resize(
sample["mask"].astype(np.float32), sample["mask"].astype(np.float32),
...@@ -149,6 +140,7 @@ class Resize(object): ...@@ -149,6 +140,7 @@ class Resize(object):
class NormalizeImage(object): class NormalizeImage(object):
"""Normlize image by given mean and std. """Normlize image by given mean and std.
""" """
def __init__(self, mean, std): def __init__(self, mean, std):
self.__mean = mean self.__mean = mean
self.__std = std self.__std = std
...@@ -162,6 +154,7 @@ class NormalizeImage(object): ...@@ -162,6 +154,7 @@ class NormalizeImage(object):
class PrepareForNet(object): class PrepareForNet(object):
"""Prepare sample for usage as network input. """Prepare sample for usage as network input.
""" """
def __init__(self): def __init__(self):
pass pass
......
...@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1): ...@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True color = True
elif (len(image.shape) == 2 elif (len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False color = False
else: else:
raise Exception( raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
"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("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0])) file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
......
...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config ...@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['InferenceModel'] __all__ = ['InferenceModel']
class InferenceModel(): class InferenceModel():
# 初始化函数 # 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True): def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True):
...@@ -27,12 +28,8 @@ class InferenceModel(): ...@@ -27,12 +28,8 @@ class InferenceModel():
''' '''
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' % ( return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (len(
len(self.input_handles), self.input_handles), str(self.input_names), len(self.output_handles), str(self.output_names))
str(self.input_names),
len(self.output_handles),
str(self.output_names)
)
# 类调用函数 # 类调用函数
def __call__(self, *input_datas, batch_size=1): def __call__(self, *input_datas, batch_size=1):
...@@ -59,7 +56,9 @@ class InferenceModel(): ...@@ -59,7 +56,9 @@ class InferenceModel():
try: try:
int(os.environ.get('CUDA_VISIBLE_DEVICES')) int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception: 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 use_gpu = False
# 加载模型参数 # 加载模型参数
......
...@@ -26,12 +26,7 @@ class MiDaS_Small(Module): ...@@ -26,12 +26,7 @@ class MiDaS_Small(Module):
model_path = os.path.join(self.directory, "model-small") model_path = os.path.join(self.directory, "model-small")
# 加载模型 # 加载模型
self.model = InferenceModel( self.model = InferenceModel(modelpath=model_path, use_gpu=use_gpu, use_mkldnn=False, combined=True)
modelpath=model_path,
use_gpu=use_gpu,
use_mkldnn=False,
combined=True
)
self.model.eval() self.model.eval()
# 数据预处理配置 # 数据预处理配置
...@@ -46,8 +41,7 @@ class MiDaS_Small(Module): ...@@ -46,8 +41,7 @@ class MiDaS_Small(Module):
resize_method="upper_bound", resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC, image_interpolation_method=cv2.INTER_CUBIC,
), ),
NormalizeImage(mean=[0.485, 0.456, 0.406], NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
std=[0.229, 0.224, 0.225]),
PrepareForNet() PrepareForNet()
]) ])
...@@ -113,12 +107,7 @@ class MiDaS_Small(Module): ...@@ -113,12 +107,7 @@ class MiDaS_Small(Module):
return outputs return outputs
# 深度估计函数 # 深度估计函数
def depth_estimation(self, def depth_estimation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
# 加载数据 # 加载数据
datas = self.load_datas(paths, images) datas = self.load_datas(paths, images)
......
...@@ -7,6 +7,7 @@ import cv2 ...@@ -7,6 +7,7 @@ import cv2
class Resize(object): class Resize(object):
"""Resize sample to given size (width, height). """Resize sample to given size (width, height).
""" """
def __init__(self, def __init__(self,
width, width,
height, height,
...@@ -51,12 +52,10 @@ class Resize(object): ...@@ -51,12 +52,10 @@ class Resize(object):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int) y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val: if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) * y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
self.__multiple_of).astype(int)
if y < min_val: if y < min_val:
y = (np.ceil(x / self.__multiple_of) * y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
self.__multiple_of).astype(int)
return y return y
...@@ -91,31 +90,24 @@ class Resize(object): ...@@ -91,31 +90,24 @@ class Resize(object):
# fit height # fit height
scale_width = scale_height scale_width = scale_height
else: else:
raise ValueError( raise ValueError(f"resize_method {self.__resize_method} not implemented")
f"resize_method {self.__resize_method} not implemented")
if self.__resize_method == "lower_bound": if self.__resize_method == "lower_bound":
new_height = self.constrain_to_multiple_of(scale_height * height, new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
min_val=self.__height) new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
elif self.__resize_method == "upper_bound": elif self.__resize_method == "upper_bound":
new_height = self.constrain_to_multiple_of(scale_height * height, new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
max_val=self.__height) new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
elif self.__resize_method == "minimal": elif self.__resize_method == "minimal":
new_height = self.constrain_to_multiple_of(scale_height * height) new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width) new_width = self.constrain_to_multiple_of(scale_width * width)
else: else:
raise ValueError( raise ValueError(f"resize_method {self.__resize_method} not implemented")
f"resize_method {self.__resize_method} not implemented")
return (new_width, new_height) return (new_width, new_height)
def __call__(self, sample): def __call__(self, sample):
width, height = self.get_size(sample["image"].shape[1], width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
sample["image"].shape[0])
# resize sample # resize sample
sample["image"] = cv2.resize( sample["image"] = cv2.resize(
...@@ -133,8 +125,7 @@ class Resize(object): ...@@ -133,8 +125,7 @@ class Resize(object):
) )
if "depth" in sample: if "depth" in sample:
sample["depth"] = cv2.resize(sample["depth"], (width, height), sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
interpolation=cv2.INTER_NEAREST)
sample["mask"] = cv2.resize( sample["mask"] = cv2.resize(
sample["mask"].astype(np.float32), sample["mask"].astype(np.float32),
...@@ -149,6 +140,7 @@ class Resize(object): ...@@ -149,6 +140,7 @@ class Resize(object):
class NormalizeImage(object): class NormalizeImage(object):
"""Normlize image by given mean and std. """Normlize image by given mean and std.
""" """
def __init__(self, mean, std): def __init__(self, mean, std):
self.__mean = mean self.__mean = mean
self.__std = std self.__std = std
...@@ -162,6 +154,7 @@ class NormalizeImage(object): ...@@ -162,6 +154,7 @@ class NormalizeImage(object):
class PrepareForNet(object): class PrepareForNet(object):
"""Prepare sample for usage as network input. """Prepare sample for usage as network input.
""" """
def __init__(self): def __init__(self):
pass pass
......
...@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1): ...@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True color = True
elif (len(image.shape) == 2 elif (len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False color = False
else: else:
raise Exception( raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
"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("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0])) 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 ...@@ -2,10 +2,11 @@ import numpy as np
import cv2 import cv2
from scipy import ndimage from scipy import ndimage
def get_normal_map(img): def get_normal_map(img):
img = img.astype(np.float) img = img.astype(np.float)
img = img / 255.0 img = img / 255.0
img = - img + 1 img = -img + 1
img[img < 0] = 0 img[img < 0] = 0
img[img > 1] = 1 img[img > 1] = 1
return img return img
...@@ -112,7 +113,7 @@ def resize_img_512(img): ...@@ -112,7 +113,7 @@ def resize_img_512(img):
def resize_img_512_3d(img): def resize_img_512_3d(img):
zeros = np.zeros((1, 3, 512, 512), dtype=np.float) 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)) return zeros.transpose((1, 2, 3, 0))
...@@ -122,7 +123,7 @@ def denoise_mat(img, i): ...@@ -122,7 +123,7 @@ def denoise_mat(img, i):
def show_active_img_and_save_denoise(img, path): def show_active_img_and_save_denoise(img, path):
mat = img.astype(np.float) mat = img.astype(np.float)
mat = - mat + 1 mat = -mat + 1
mat = mat * 255.0 mat = mat * 255.0
mat[mat < 0] = 0 mat[mat < 0] = 0
mat[mat > 255] = 255 mat[mat > 255] = 255
...@@ -132,11 +133,9 @@ def show_active_img_and_save_denoise(img, path): ...@@ -132,11 +133,9 @@ def show_active_img_and_save_denoise(img, path):
return return
def show_active_img(name, img): def show_active_img(name, img):
mat = img.astype(np.float) mat = img.astype(np.float)
mat = - mat + 1 mat = -mat + 1
mat = mat * 255.0 mat = mat * 255.0
mat[mat < 0] = 0 mat[mat < 0] = 0
mat[mat > 255] = 255 mat[mat > 255] = 255
...@@ -147,7 +146,7 @@ def show_active_img(name, img): ...@@ -147,7 +146,7 @@ def show_active_img(name, img):
def get_active_img(img): def get_active_img(img):
mat = img.astype(np.float) mat = img.astype(np.float)
mat = - mat + 1 mat = -mat + 1
mat = mat * 255.0 mat = mat * 255.0
mat[mat < 0] = 0 mat[mat < 0] = 0
mat[mat > 255] = 255 mat[mat > 255] = 255
...@@ -158,7 +157,7 @@ def get_active_img(img): ...@@ -158,7 +157,7 @@ def get_active_img(img):
def get_active_img_fil(img): def get_active_img_fil(img):
mat = img.astype(np.float) mat = img.astype(np.float)
mat[mat < 0.18] = 0 mat[mat < 0.18] = 0
mat = - mat + 1 mat = -mat + 1
mat = mat * 255.0 mat = mat * 255.0
mat[mat < 0] = 0 mat[mat < 0] = 0
mat[mat > 255] = 255 mat[mat > 255] = 255
......
...@@ -14,12 +14,12 @@ import paddle.fluid as fluid ...@@ -14,12 +14,12 @@ import paddle.fluid as fluid
import paddlehub as hub import paddlehub as hub
from Extract_Line_Draft.function import * from Extract_Line_Draft.function import *
@moduleinfo( @moduleinfo(
name="Extract_Line_Draft", name="Extract_Line_Draft",
version="1.0.0", version="1.0.0",
type="cv/segmentation", type="cv/segmentation",
summary= summary="Import the color picture and generate the line draft of the picture",
"Import the color picture and generate the line draft of the picture",
author="彭兆帅,郑博培", author="彭兆帅,郑博培",
author_email="1084667371@qq.com,2733821739@qq.com") author_email="1084667371@qq.com,2733821739@qq.com")
class ExtractLineDraft(hub.Module): class ExtractLineDraft(hub.Module):
...@@ -28,8 +28,9 @@ class ExtractLineDraft(hub.Module): ...@@ -28,8 +28,9 @@ class ExtractLineDraft(hub.Module):
Initialize with the necessary elements 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() self._set_config()
def _set_config(self): def _set_config(self):
""" """
predictor config setting predictor config setting
...@@ -153,12 +154,9 @@ class ExtractLineDraft(hub.Module): ...@@ -153,12 +154,9 @@ class ExtractLineDraft(hub.Module):
usage='%(prog)s', usage='%(prog)s',
add_help=True) add_help=True)
self.arg_input_group = self.parser.add_argument_group( self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required")
title="Input options", description="Input data. Required")
self.arg_config_group = self.parser.add_argument_group( self.arg_config_group = self.parser.add_argument_group(
title="Config options", title="Config options", description="Run configuration for controlling module behavior, not required.")
description=
"Run configuration for controlling module behavior, not required.")
self.add_module_input_arg() self.add_module_input_arg()
...@@ -173,21 +171,12 @@ class ExtractLineDraft(hub.Module): ...@@ -173,21 +171,12 @@ class ExtractLineDraft(hub.Module):
use_gpu = args.use_gpu use_gpu = args.use_gpu
self.ExtractLine(image=input_data, use_gpu=use_gpu) self.ExtractLine(image=input_data, use_gpu=use_gpu)
def add_module_input_arg(self): def add_module_input_arg(self):
""" """
Add the command input options Add the command input options
""" """
self.arg_input_group.add_argument( self.arg_input_group.add_argument('--image', type=str, default=None, help="file contain input data")
'--image', self.arg_input_group.add_argument('--use_gpu', type=ast.literal_eval, default=None, help="weather to use gpu")
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): def check_input_data(self, args):
input_data = [] input_data = []
...@@ -196,4 +185,3 @@ class ExtractLineDraft(hub.Module): ...@@ -196,4 +185,3 @@ class ExtractLineDraft(hub.Module):
raise RuntimeError("Path %s is not exist." % args.image) raise RuntimeError("Path %s is not exist." % args.image)
path = "{}".format(args.image) path = "{}".format(args.image)
return path return path
...@@ -48,12 +48,9 @@ class FCN(nn.Layer): ...@@ -48,12 +48,9 @@ class FCN(nn.Layer):
super(FCN, self).__init__() super(FCN, self).__init__()
self.backbone = backbone self.backbone = backbone
backbone_channels = [ backbone_channels = [backbone.feat_channels[i] for i in backbone_indices]
backbone.feat_channels[i] for i in backbone_indices
]
self.head = FCNHead(num_classes, backbone_indices, backbone_channels, self.head = FCNHead(num_classes, backbone_indices, backbone_channels, channels)
channels)
self.align_corners = align_corners self.align_corners = align_corners
self.pretrained = pretrained self.pretrained = pretrained
...@@ -62,11 +59,7 @@ class FCN(nn.Layer): ...@@ -62,11 +59,7 @@ class FCN(nn.Layer):
feat_list = self.backbone(x) feat_list = self.backbone(x)
logit_list = self.head(feat_list) logit_list = self.head(feat_list)
return [ return [
F.interpolate( F.interpolate(logit, x.shape[2:], mode='bilinear', align_corners=self.align_corners) for logit in logit_list
logit,
x.shape[2:],
mode='bilinear',
align_corners=self.align_corners) for logit in logit_list
] ]
...@@ -83,11 +76,7 @@ class FCNHead(nn.Layer): ...@@ -83,11 +76,7 @@ class FCNHead(nn.Layer):
pretrained (str, optional): The path of pretrained model. Default: None pretrained (str, optional): The path of pretrained model. Default: None
""" """
def __init__(self, def __init__(self, num_classes, backbone_indices=(-1, ), backbone_channels=(270, ), channels=None):
num_classes,
backbone_indices=(-1, ),
backbone_channels=(270, ),
channels=None):
super(FCNHead, self).__init__() super(FCNHead, self).__init__()
self.num_classes = num_classes self.num_classes = num_classes
...@@ -96,17 +85,8 @@ class FCNHead(nn.Layer): ...@@ -96,17 +85,8 @@ class FCNHead(nn.Layer):
channels = backbone_channels[0] channels = backbone_channels[0]
self.conv_1 = ConvBNReLU( self.conv_1 = ConvBNReLU(
in_channels=backbone_channels[0], in_channels=backbone_channels[0], out_channels=channels, kernel_size=1, padding='same', stride=1)
out_channels=channels, self.cls = nn.Conv2D(in_channels=channels, out_channels=self.num_classes, kernel_size=1, stride=1, padding=0)
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): def forward(self, feat_list):
logit_list = [] logit_list = []
...@@ -115,4 +95,3 @@ class FCNHead(nn.Layer): ...@@ -115,4 +95,3 @@ class FCNHead(nn.Layer):
logit = self.cls(x) logit = self.cls(x)
logit_list.append(logit) logit_list.append(logit)
return logit_list return logit_list
...@@ -21,8 +21,8 @@ import paddle.nn.functional as F ...@@ -21,8 +21,8 @@ import paddle.nn.functional as F
from .layers import ConvBNReLU, ConvBN from .layers import ConvBNReLU, ConvBN
__all__ = [ __all__ = [
"HRNet_W18_Small_V1", "HRNet_W18_Small_V2", "HRNet_W18", "HRNet_W30", "HRNet_W18_Small_V1", "HRNet_W18_Small_V2", "HRNet_W18", "HRNet_W30", "HRNet_W32", "HRNet_W40", "HRNet_W44",
"HRNet_W32", "HRNet_W40", "HRNet_W44", "HRNet_W48", "HRNet_W60", "HRNet_W64" "HRNet_W48", "HRNet_W60", "HRNet_W64"
] ]
...@@ -88,20 +88,10 @@ class HRNet(nn.Layer): ...@@ -88,20 +88,10 @@ class HRNet(nn.Layer):
self.feat_channels = [sum(stage4_num_channels)] self.feat_channels = [sum(stage4_num_channels)]
self.conv_layer1_1 = ConvBNReLU( self.conv_layer1_1 = ConvBNReLU(
in_channels=3, in_channels=3, out_channels=64, kernel_size=3, stride=2, padding='same', bias_attr=False)
out_channels=64,
kernel_size=3,
stride=2,
padding='same',
bias_attr=False)
self.conv_layer1_2 = ConvBNReLU( self.conv_layer1_2 = ConvBNReLU(
in_channels=64, in_channels=64, out_channels=64, kernel_size=3, stride=2, padding='same', bias_attr=False)
out_channels=64,
kernel_size=3,
stride=2,
padding='same',
bias_attr=False)
self.la1 = Layer1( self.la1 = Layer1(
num_channels=64, num_channels=64,
...@@ -111,9 +101,7 @@ class HRNet(nn.Layer): ...@@ -111,9 +101,7 @@ class HRNet(nn.Layer):
name="layer2") name="layer2")
self.tr1 = TransitionLayer( self.tr1 = TransitionLayer(
in_channels=[self.stage1_num_channels[0] * 4], in_channels=[self.stage1_num_channels[0] * 4], out_channels=self.stage2_num_channels, name="tr1")
out_channels=self.stage2_num_channels,
name="tr1")
self.st2 = Stage( self.st2 = Stage(
num_channels=self.stage2_num_channels, num_channels=self.stage2_num_channels,
...@@ -125,9 +113,7 @@ class HRNet(nn.Layer): ...@@ -125,9 +113,7 @@ class HRNet(nn.Layer):
align_corners=align_corners) align_corners=align_corners)
self.tr2 = TransitionLayer( self.tr2 = TransitionLayer(
in_channels=self.stage2_num_channels, in_channels=self.stage2_num_channels, out_channels=self.stage3_num_channels, name="tr2")
out_channels=self.stage3_num_channels,
name="tr2")
self.st3 = Stage( self.st3 = Stage(
num_channels=self.stage3_num_channels, num_channels=self.stage3_num_channels,
num_modules=self.stage3_num_modules, num_modules=self.stage3_num_modules,
...@@ -138,9 +124,7 @@ class HRNet(nn.Layer): ...@@ -138,9 +124,7 @@ class HRNet(nn.Layer):
align_corners=align_corners) align_corners=align_corners)
self.tr3 = TransitionLayer( self.tr3 = TransitionLayer(
in_channels=self.stage3_num_channels, in_channels=self.stage3_num_channels, out_channels=self.stage4_num_channels, name="tr3")
out_channels=self.stage4_num_channels,
name="tr3")
self.st4 = Stage( self.st4 = Stage(
num_channels=self.stage4_num_channels, num_channels=self.stage4_num_channels,
num_modules=self.stage4_num_modules, num_modules=self.stage4_num_modules,
...@@ -166,30 +150,16 @@ class HRNet(nn.Layer): ...@@ -166,30 +150,16 @@ class HRNet(nn.Layer):
st4 = self.st4(tr3) st4 = self.st4(tr3)
x0_h, x0_w = st4[0].shape[2:] x0_h, x0_w = st4[0].shape[2:]
x1 = F.interpolate( x1 = F.interpolate(st4[1], (x0_h, x0_w), mode='bilinear', align_corners=self.align_corners)
st4[1], (x0_h, x0_w), x2 = F.interpolate(st4[2], (x0_h, x0_w), mode='bilinear', align_corners=self.align_corners)
mode='bilinear', x3 = F.interpolate(st4[3], (x0_h, x0_w), mode='bilinear', align_corners=self.align_corners)
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) x = paddle.concat([st4[0], x1, x2, x3], axis=1)
return [x] return [x]
class Layer1(nn.Layer): class Layer1(nn.Layer):
def __init__(self, def __init__(self, num_channels, num_filters, num_blocks, has_se=False, name=None):
num_channels,
num_filters,
num_blocks,
has_se=False,
name=None):
super(Layer1, self).__init__() super(Layer1, self).__init__()
self.bottleneck_block_list = [] self.bottleneck_block_list = []
...@@ -258,12 +228,7 @@ class TransitionLayer(nn.Layer): ...@@ -258,12 +228,7 @@ class TransitionLayer(nn.Layer):
class Branches(nn.Layer): class Branches(nn.Layer):
def __init__(self, def __init__(self, num_blocks, in_channels, out_channels, has_se=False, name=None):
num_blocks,
in_channels,
out_channels,
has_se=False,
name=None):
super(Branches, self).__init__() super(Branches, self).__init__()
self.basic_block_list = [] self.basic_block_list = []
...@@ -278,8 +243,7 @@ class Branches(nn.Layer): ...@@ -278,8 +243,7 @@ class Branches(nn.Layer):
num_channels=in_ch, num_channels=in_ch,
num_filters=out_channels[i], num_filters=out_channels[i],
has_se=has_se, has_se=has_se,
name=name + '_branch_layer_' + str(i + 1) + '_' + name=name + '_branch_layer_' + str(i + 1) + '_' + str(j + 1)))
str(j + 1)))
self.basic_block_list[i].append(basic_block_func) self.basic_block_list[i].append(basic_block_func)
def forward(self, x): def forward(self, x):
...@@ -293,24 +257,14 @@ class Branches(nn.Layer): ...@@ -293,24 +257,14 @@ class Branches(nn.Layer):
class BottleneckBlock(nn.Layer): class BottleneckBlock(nn.Layer):
def __init__(self, def __init__(self, num_channels, num_filters, has_se, stride=1, downsample=False, name=None):
num_channels,
num_filters,
has_se,
stride=1,
downsample=False,
name=None):
super(BottleneckBlock, self).__init__() super(BottleneckBlock, self).__init__()
self.has_se = has_se self.has_se = has_se
self.downsample = downsample self.downsample = downsample
self.conv1 = ConvBNReLU( self.conv1 = ConvBNReLU(
in_channels=num_channels, in_channels=num_channels, out_channels=num_filters, kernel_size=1, padding='same', bias_attr=False)
out_channels=num_filters,
kernel_size=1,
padding='same',
bias_attr=False)
self.conv2 = ConvBNReLU( self.conv2 = ConvBNReLU(
in_channels=num_filters, in_channels=num_filters,
...@@ -321,26 +275,15 @@ class BottleneckBlock(nn.Layer): ...@@ -321,26 +275,15 @@ class BottleneckBlock(nn.Layer):
bias_attr=False) bias_attr=False)
self.conv3 = ConvBN( self.conv3 = ConvBN(
in_channels=num_filters, in_channels=num_filters, out_channels=num_filters * 4, kernel_size=1, padding='same', bias_attr=False)
out_channels=num_filters * 4,
kernel_size=1,
padding='same',
bias_attr=False)
if self.downsample: if self.downsample:
self.conv_down = ConvBN( self.conv_down = ConvBN(
in_channels=num_channels, in_channels=num_channels, out_channels=num_filters * 4, kernel_size=1, padding='same', bias_attr=False)
out_channels=num_filters * 4,
kernel_size=1,
padding='same',
bias_attr=False)
if self.has_se: if self.has_se:
self.se = SELayer( self.se = SELayer(
num_channels=num_filters * 4, num_channels=num_filters * 4, num_filters=num_filters * 4, reduction_ratio=16, name=name + '_fc')
num_filters=num_filters * 4,
reduction_ratio=16,
name=name + '_fc')
def forward(self, x): def forward(self, x):
residual = x residual = x
...@@ -360,13 +303,7 @@ class BottleneckBlock(nn.Layer): ...@@ -360,13 +303,7 @@ class BottleneckBlock(nn.Layer):
class BasicBlock(nn.Layer): class BasicBlock(nn.Layer):
def __init__(self, def __init__(self, num_channels, num_filters, stride=1, has_se=False, downsample=False, name=None):
num_channels,
num_filters,
stride=1,
has_se=False,
downsample=False,
name=None):
super(BasicBlock, self).__init__() super(BasicBlock, self).__init__()
self.has_se = has_se self.has_se = has_se
...@@ -380,26 +317,14 @@ class BasicBlock(nn.Layer): ...@@ -380,26 +317,14 @@ class BasicBlock(nn.Layer):
padding='same', padding='same',
bias_attr=False) bias_attr=False)
self.conv2 = ConvBN( self.conv2 = ConvBN(
in_channels=num_filters, in_channels=num_filters, out_channels=num_filters, kernel_size=3, padding='same', bias_attr=False)
out_channels=num_filters,
kernel_size=3,
padding='same',
bias_attr=False)
if self.downsample: if self.downsample:
self.conv_down = ConvBNReLU( self.conv_down = ConvBNReLU(
in_channels=num_channels, in_channels=num_channels, out_channels=num_filters, kernel_size=1, padding='same', bias_attr=False)
out_channels=num_filters,
kernel_size=1,
padding='same',
bias_attr=False)
if self.has_se: if self.has_se:
self.se = SELayer( self.se = SELayer(num_channels=num_filters, num_filters=num_filters, reduction_ratio=16, name=name + '_fc')
num_channels=num_filters,
num_filters=num_filters,
reduction_ratio=16,
name=name + '_fc')
def forward(self, x): def forward(self, x):
residual = x residual = x
...@@ -428,17 +353,11 @@ class SELayer(nn.Layer): ...@@ -428,17 +353,11 @@ class SELayer(nn.Layer):
med_ch = int(num_channels / reduction_ratio) med_ch = int(num_channels / reduction_ratio)
stdv = 1.0 / math.sqrt(num_channels * 1.0) stdv = 1.0 / math.sqrt(num_channels * 1.0)
self.squeeze = nn.Linear( self.squeeze = nn.Linear(
num_channels, num_channels, med_ch, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Uniform(-stdv, stdv)))
med_ch,
weight_attr=paddle.ParamAttr(
initializer=nn.initializer.Uniform(-stdv, stdv)))
stdv = 1.0 / math.sqrt(med_ch * 1.0) stdv = 1.0 / math.sqrt(med_ch * 1.0)
self.excitation = nn.Linear( self.excitation = nn.Linear(
med_ch, med_ch, num_filters, weight_attr=paddle.ParamAttr(initializer=nn.initializer.Uniform(-stdv, stdv)))
num_filters,
weight_attr=paddle.ParamAttr(
initializer=nn.initializer.Uniform(-stdv, stdv)))
def forward(self, x): def forward(self, x):
pool = self.pool2d_gap(x) pool = self.pool2d_gap(x)
...@@ -447,8 +366,7 @@ class SELayer(nn.Layer): ...@@ -447,8 +366,7 @@ class SELayer(nn.Layer):
squeeze = F.relu(squeeze) squeeze = F.relu(squeeze)
excitation = self.excitation(squeeze) excitation = self.excitation(squeeze)
excitation = F.sigmoid(excitation) excitation = F.sigmoid(excitation)
excitation = paddle.reshape( excitation = paddle.reshape(excitation, shape=[-1, self._num_channels, 1, 1])
excitation, shape=[-1, self._num_channels, 1, 1])
out = x * excitation out = x * excitation
return out return out
...@@ -512,11 +430,7 @@ class HighResolutionModule(nn.Layer): ...@@ -512,11 +430,7 @@ class HighResolutionModule(nn.Layer):
super(HighResolutionModule, self).__init__() super(HighResolutionModule, self).__init__()
self.branches_func = Branches( self.branches_func = Branches(
num_blocks=num_blocks, num_blocks=num_blocks, in_channels=num_channels, out_channels=num_filters, has_se=has_se, name=name)
in_channels=num_channels,
out_channels=num_filters,
has_se=has_se,
name=name)
self.fuse_func = FuseLayers( self.fuse_func = FuseLayers(
in_channels=num_filters, in_channels=num_filters,
...@@ -532,12 +446,7 @@ class HighResolutionModule(nn.Layer): ...@@ -532,12 +446,7 @@ class HighResolutionModule(nn.Layer):
class FuseLayers(nn.Layer): class FuseLayers(nn.Layer):
def __init__(self, def __init__(self, in_channels, out_channels, multi_scale_output=True, name=None, align_corners=False):
in_channels,
out_channels,
multi_scale_output=True,
name=None,
align_corners=False):
super(FuseLayers, self).__init__() super(FuseLayers, self).__init__()
self._actual_ch = len(in_channels) if multi_scale_output else 1 self._actual_ch = len(in_channels) if multi_scale_output else 1
...@@ -562,8 +471,7 @@ class FuseLayers(nn.Layer): ...@@ -562,8 +471,7 @@ class FuseLayers(nn.Layer):
for k in range(i - j): for k in range(i - j):
if k == i - j - 1: if k == i - j - 1:
residual_func = self.add_sublayer( residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}_{}".format( "residual_{}_layer_{}_{}_{}".format(name, i + 1, j + 1, k + 1),
name, i + 1, j + 1, k + 1),
ConvBN( ConvBN(
in_channels=pre_num_filters, in_channels=pre_num_filters,
out_channels=out_channels[i], out_channels=out_channels[i],
...@@ -574,8 +482,7 @@ class FuseLayers(nn.Layer): ...@@ -574,8 +482,7 @@ class FuseLayers(nn.Layer):
pre_num_filters = out_channels[i] pre_num_filters = out_channels[i]
else: else:
residual_func = self.add_sublayer( residual_func = self.add_sublayer(
"residual_{}_layer_{}_{}_{}".format( "residual_{}_layer_{}_{}_{}".format(name, i + 1, j + 1, k + 1),
name, i + 1, j + 1, k + 1),
ConvBNReLU( ConvBNReLU(
in_channels=pre_num_filters, in_channels=pre_num_filters,
out_channels=out_channels[j], out_channels=out_channels[j],
...@@ -597,11 +504,7 @@ class FuseLayers(nn.Layer): ...@@ -597,11 +504,7 @@ class FuseLayers(nn.Layer):
y = self.residual_func_list[residual_func_idx](x[j]) y = self.residual_func_list[residual_func_idx](x[j])
residual_func_idx += 1 residual_func_idx += 1
y = F.interpolate( y = F.interpolate(y, residual_shape, mode='bilinear', align_corners=self.align_corners)
y,
residual_shape,
mode='bilinear',
align_corners=self.align_corners)
residual = residual + y residual = residual + y
elif j < i: elif j < i:
y = x[j] y = x[j]
......
...@@ -26,16 +26,10 @@ def SyncBatchNorm(*args, **kwargs): ...@@ -26,16 +26,10 @@ def SyncBatchNorm(*args, **kwargs):
class ConvBNReLU(nn.Layer): class ConvBNReLU(nn.Layer):
def __init__(self, def __init__(self, in_channels, out_channels, kernel_size, padding='same', **kwargs):
in_channels,
out_channels,
kernel_size,
padding='same',
**kwargs):
super().__init__() super().__init__()
self._conv = nn.Conv2D( self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding, **kwargs)
in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._batch_norm = SyncBatchNorm(out_channels) self._batch_norm = SyncBatchNorm(out_channels)
...@@ -47,19 +41,12 @@ class ConvBNReLU(nn.Layer): ...@@ -47,19 +41,12 @@ class ConvBNReLU(nn.Layer):
class ConvBN(nn.Layer): class ConvBN(nn.Layer):
def __init__(self, def __init__(self, in_channels, out_channels, kernel_size, padding='same', **kwargs):
in_channels,
out_channels,
kernel_size,
padding='same',
**kwargs):
super().__init__() super().__init__()
self._conv = nn.Conv2D( self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding, **kwargs)
in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._batch_norm = SyncBatchNorm(out_channels) self._batch_norm = SyncBatchNorm(out_channels)
def forward(self, x): def forward(self, x):
x = self._conv(x) x = self._conv(x)
x = self._batch_norm(x) x = self._batch_norm(x)
return x return x
此差异已折叠。
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册