提交 d2444b7e 编写于 作者: W wuzewu

Add thirdparty directory

上级 7cad8e4e
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self,in_ch=3,out_ch=3,dirate=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2D(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer):#UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer):#UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer):#UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer):#UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer):#UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx2d = self.rebnconv2d(paddle.concat((hx3d,hx2),1))
hx1d = self.rebnconv1d(paddle.concat((hx2d,hx1),1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NET,self).__init__()
self.stage1 = RSU7(in_ch,32,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(128,out_ch,3,padding=1)
self.side4 = nn.Conv2D(256,out_ch,3,padding=1)
self.side5 = nn.Conv2D(512,out_ch,3,padding=1)
self.side6 = nn.Conv2D(512,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NETP,self).__init__()
self.stage1 = RSU7(in_ch,16,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,16,64)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(64,16,64)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(64,16,64)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(64,16,64)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(64,16,64)
# decoder
self.stage5d = RSU4F(128,16,64)
self.stage4d = RSU4(128,16,64)
self.stage3d = RSU5(128,16,64)
self.stage2d = RSU6(128,16,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(64,out_ch,3,padding=1)
self.side4 = nn.Conv2D(64,out_ch,3,padding=1)
self.side5 = nn.Conv2D(64,out_ch,3,padding=1)
self.side6 = nn.Conv2D(64,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self,in_ch=3,out_ch=3,dirate=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2D(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer):#UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer):#UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer):#UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer):#UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer):#UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx2d = self.rebnconv2d(paddle.concat((hx3d,hx2),1))
hx1d = self.rebnconv1d(paddle.concat((hx2d,hx1),1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NET,self).__init__()
self.stage1 = RSU7(in_ch,32,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(128,out_ch,3,padding=1)
self.side4 = nn.Conv2D(256,out_ch,3,padding=1)
self.side5 = nn.Conv2D(512,out_ch,3,padding=1)
self.side6 = nn.Conv2D(512,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NETP,self).__init__()
self.stage1 = RSU7(in_ch,16,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,16,64)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(64,16,64)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(64,16,64)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(64,16,64)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(64,16,64)
# decoder
self.stage5d = RSU4F(128,16,64)
self.stage4d = RSU4(128,16,64)
self.stage3d = RSU5(128,16,64)
self.stage2d = RSU6(128,16,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(64,out_ch,3,padding=1)
self.side4 = nn.Conv2D(64,out_ch,3,padding=1)
self.side5 = nn.Conv2D(64,out_ch,3,padding=1)
self.side6 = nn.Conv2D(64,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self,in_ch=3,out_ch=3,dirate=1):
super(REBNCONV,self).__init__()
self.conv_s1 = nn.Conv2D(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self,x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src,tar):
src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer):#UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool5 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7,hx6),1))
hx6dup = _upsample_like(hx6d,hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer):#UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool4 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer):#UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool3 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer):#UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.pool1 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.pool2 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup,hx1),1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer):#UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F,self).__init__()
self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
def forward(self,x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4,hx3),1))
hx2d = self.rebnconv2d(paddle.concat((hx3d,hx2),1))
hx1d = self.rebnconv1d(paddle.concat((hx2d,hx1),1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NET,self).__init__()
self.stage1 = RSU7(in_ch,32,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,32,128)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(128,64,256)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(256,128,512)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(512,256,512)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(512,256,512)
# decoder
self.stage5d = RSU4F(1024,256,512)
self.stage4d = RSU4(1024,128,256)
self.stage3d = RSU5(512,64,128)
self.stage2d = RSU6(256,32,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(128,out_ch,3,padding=1)
self.side4 = nn.Conv2D(256,out_ch,3,padding=1)
self.side5 = nn.Conv2D(512,out_ch,3,padding=1)
self.side6 = nn.Conv2D(512,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self,in_ch=3,out_ch=1):
super(U2NETP,self).__init__()
self.stage1 = RSU7(in_ch,16,64)
self.pool12 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage2 = RSU6(64,16,64)
self.pool23 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage3 = RSU5(64,16,64)
self.pool34 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage4 = RSU4(64,16,64)
self.pool45 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage5 = RSU4F(64,16,64)
self.pool56 = nn.MaxPool2D(2,stride=2,ceil_mode=True)
self.stage6 = RSU4F(64,16,64)
# decoder
self.stage5d = RSU4F(128,16,64)
self.stage4d = RSU4(128,16,64)
self.stage3d = RSU5(128,16,64)
self.stage2d = RSU6(128,16,64)
self.stage1d = RSU7(128,16,64)
self.side1 = nn.Conv2D(64,out_ch,3,padding=1)
self.side2 = nn.Conv2D(64,out_ch,3,padding=1)
self.side3 = nn.Conv2D(64,out_ch,3,padding=1)
self.side4 = nn.Conv2D(64,out_ch,3,padding=1)
self.side5 = nn.Conv2D(64,out_ch,3,padding=1)
self.side6 = nn.Conv2D(64,out_ch,3,padding=1)
self.outconv = nn.Conv2D(6,out_ch,1)
def forward(self,x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6,hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up,hx5),1))
hx5dup = _upsample_like(hx5d,hx4)
hx4d = self.stage4d(paddle.concat((hx5dup,hx4),1))
hx4dup = _upsample_like(hx4d,hx3)
hx3d = self.stage3d(paddle.concat((hx4dup,hx3),1))
hx3dup = _upsample_like(hx3d,hx2)
hx2d = self.stage2d(paddle.concat((hx3dup,hx2),1))
hx2dup = _upsample_like(hx2d,hx1)
hx1d = self.stage1d(paddle.concat((hx2dup,hx1),1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2,d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3,d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4,d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5,d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6,d1)
d0 = self.outconv(paddle.concat((d1,d2,d3,d4,d5,d6),1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
......@@ -7,13 +7,14 @@ import paddle.nn as nn
import paddlehub as hub
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="ID_Photo_GEN", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="ID_Photo_GEN", # 模型介绍
version="1.0.0" # 版本号
name="ID_Photo_GEN", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="ID_Photo_GEN", # 模型介绍
version="1.0.0" # 版本号
)
class ID_Photo_GEN(nn.Layer):
def __init__(self):
......@@ -38,18 +39,15 @@ class ID_Photo_GEN(nn.Layer):
if images is not None:
datas = images
# 返回数据列表
return datas
# 数据预处理函数
def preprocess(self, images, batch_size, use_gpu):
# 获取人脸关键点
outputs = self.face_detector.keypoint_detection(
images=images,
batch_size=batch_size,
use_gpu=use_gpu)
outputs = self.face_detector.keypoint_detection(images=images, batch_size=batch_size, use_gpu=use_gpu)
crops = []
for output, image in zip(outputs, images):
for landmarks in output['data']:
......@@ -58,7 +56,8 @@ class ID_Photo_GEN(nn.Layer):
# rotation angle
left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45]
radian = np.arctan((left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
radian = np.arctan(
(left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
# image size after rotating
height, width, _ = image.shape
......@@ -73,7 +72,7 @@ class ID_Photo_GEN(nn.Layer):
# affine matrix
M = np.array([[cos, sin, (1 - cos) * width / 2. - sin * height / 2. + Tx],
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
image = cv2.warpAffine(image, M, (new_w, new_h), borderValue=(255, 255, 255))
......@@ -103,16 +102,17 @@ class ID_Photo_GEN(nn.Layer):
h, w = image.shape[:2]
left_white = max(0, -left)
left = max(0, left)
right = min(right, w-1)
right_white = left_white + (right-left)
right = min(right, w - 1)
right_white = left_white + (right - left)
top_white = max(0, -top)
top = max(0, top)
bottom = min(bottom, h-1)
bottom = min(bottom, h - 1)
bottom_white = top_white + (bottom - top)
image_crop[top_white:bottom_white+1, left_white:right_white+1] = image[top:bottom+1, left:right+1].copy()
image_crop[top_white:bottom_white + 1, left_white:right_white + 1] = image[top:bottom + 1, left:right +
1].copy()
crops.append(image_crop)
# 获取人像分割的输出
results = self.seg.Segmentation(images=crops, batch_size=batch_size)
......@@ -141,7 +141,7 @@ class ID_Photo_GEN(nn.Layer):
cartoon = self.net(data)
outputs.append(cartoon[0].numpy())
outputs = np.concatenate(outputs, 0)
return outputs
......@@ -168,22 +168,11 @@ class ID_Photo_GEN(nn.Layer):
cv2.imwrite(os.path.join(output_dir, 'blue_%d.jpg' % i), blue)
cv2.imwrite(os.path.join(output_dir, 'red_%d.jpg' % i), red)
results.append({
'write': write,
'blue': blue,
'red': red
})
results.append({'write': write, 'blue': blue, 'red': red})
return results
def Photo_GEN(
self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
def Photo_GEN(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False, use_gpu=False):
# 获取输入数据
images = self.load_datas(paths, images)
......@@ -194,4 +183,4 @@ class ID_Photo_GEN(nn.Layer):
# 结果后处理
results = self.postprocess(faces, masks, visualization, output_dir)
return results
\ No newline at end of file
return results
......@@ -17,72 +17,66 @@ class ResnetGenerator(nn.Layer):
nn.ReLU()
]
DownBlock += [
HourGlass(ngf, ngf),
HourGlass(ngf, ngf)
]
DownBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf)]
# Down-Sampling
n_downsampling = 2
for i in range(n_downsampling):
mult = 2 ** i
mult = 2**i
DownBlock += [
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf*mult, ngf*mult*2, kernel_size=3, stride=2, bias_attr=False),
nn.InstanceNorm2D(ngf*mult*2, weight_attr=False, bias_attr=False),
nn.Conv2D(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, bias_attr=False),
nn.InstanceNorm2D(ngf * mult * 2, weight_attr=False, bias_attr=False),
nn.ReLU()
]
# Encoder Bottleneck
mult = 2 ** n_downsampling
mult = 2**n_downsampling
for i in range(n_blocks):
setattr(self, 'EncodeBlock'+str(i+1), ResnetBlock(ngf*mult))
setattr(self, 'EncodeBlock' + str(i + 1), ResnetBlock(ngf * mult))
# Class Activation Map
self.gap_fc = nn.Linear(ngf*mult, 1, bias_attr=False)
self.gmp_fc = nn.Linear(ngf*mult, 1, bias_attr=False)
self.conv1x1 = nn.Conv2D(ngf*mult*2, ngf*mult, kernel_size=1, stride=1)
self.gap_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.gmp_fc = nn.Linear(ngf * mult, 1, bias_attr=False)
self.conv1x1 = nn.Conv2D(ngf * mult * 2, ngf * mult, kernel_size=1, stride=1)
self.relu = nn.ReLU()
# Gamma, Beta block
FC = []
if self.light:
FC += [
nn.Linear(ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(),
nn.Linear(ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
else:
FC += [
nn.Linear(img_size//mult*img_size//mult*ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(img_size // mult * img_size // mult * ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU(),
nn.Linear(ngf*mult, ngf*mult, bias_attr=False),
nn.Linear(ngf * mult, ngf * mult, bias_attr=False),
nn.ReLU()
]
# Decoder Bottleneck
mult = 2 ** n_downsampling
mult = 2**n_downsampling
for i in range(n_blocks):
setattr(self, 'DecodeBlock'+str(i + 1), ResnetSoftAdaLINBlock(ngf*mult))
setattr(self, 'DecodeBlock' + str(i + 1), ResnetSoftAdaLINBlock(ngf * mult))
# Up-Sampling
UpBlock = []
for i in range(n_downsampling):
mult = 2 ** (n_downsampling - i)
mult = 2**(n_downsampling - i)
UpBlock += [
nn.Upsample(scale_factor=2),
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(ngf*mult, ngf*mult//2, kernel_size=3, stride=1, bias_attr=False),
LIN(ngf*mult//2),
nn.Conv2D(ngf * mult, ngf * mult // 2, kernel_size=3, stride=1, bias_attr=False),
LIN(ngf * mult // 2),
nn.ReLU()
]
UpBlock += [
HourGlass(ngf, ngf),
HourGlass(ngf, ngf, False)
]
UpBlock += [HourGlass(ngf, ngf), HourGlass(ngf, ngf, False)]
UpBlock += [
nn.Pad2D([3, 3, 3, 3], 'reflect'),
......@@ -101,7 +95,7 @@ class ResnetGenerator(nn.Layer):
content_features = []
for i in range(self.n_blocks):
x = getattr(self, 'EncodeBlock'+str(i+1))(x)
x = getattr(self, 'EncodeBlock' + str(i + 1))(x)
content_features.append(F.adaptive_avg_pool2d(x, 1).reshape([bs, -1]))
gap = F.adaptive_avg_pool2d(x, 1)
......@@ -127,7 +121,7 @@ class ResnetGenerator(nn.Layer):
style_features = self.FC(x.reshape([bs, -1]))
for i in range(self.n_blocks):
x = getattr(self, 'DecodeBlock'+str(i+1))(x, content_features[4-i-1], style_features)
x = getattr(self, 'DecodeBlock' + str(i + 1))(x, content_features[4 - i - 1], style_features)
out = self.UpBlock(x)
......@@ -140,25 +134,20 @@ class ConvBlock(nn.Layer):
self.dim_in = dim_in
self.dim_out = dim_out
self.conv_block1 = self.__convblock(dim_in, dim_out//2)
self.conv_block2 = self.__convblock(dim_out//2, dim_out//4)
self.conv_block3 = self.__convblock(dim_out//4, dim_out//4)
self.conv_block1 = self.__convblock(dim_in, dim_out // 2)
self.conv_block2 = self.__convblock(dim_out // 2, dim_out // 4)
self.conv_block3 = self.__convblock(dim_out // 4, dim_out // 4)
if self.dim_in != self.dim_out:
self.conv_skip = nn.Sequential(
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False),
nn.ReLU(),
nn.Conv2D(dim_in, dim_out, kernel_size=1, stride=1, bias_attr=False)
)
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.ReLU(),
nn.Conv2D(dim_in, dim_out, kernel_size=1, stride=1, bias_attr=False))
@staticmethod
def __convblock(dim_in, dim_out):
return nn.Sequential(
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False),
nn.ReLU(),
nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(dim_in, dim_out, kernel_size=3, stride=1, bias_attr=False)
)
nn.InstanceNorm2D(dim_in, weight_attr=False, bias_attr=False), nn.ReLU(), nn.Pad2D([1, 1, 1, 1], 'reflect'),
nn.Conv2D(dim_in, dim_out, kernel_size=3, stride=1, bias_attr=False))
def forward(self, x):
residual = x
......@@ -182,24 +171,24 @@ class HourGlassBlock(nn.Layer):
self.n_block = 9
for i in range(self.n_skip):
setattr(self, 'ConvBlockskip'+str(i+1), ConvBlock(dim_in, dim_in))
setattr(self, 'ConvBlockskip' + str(i + 1), ConvBlock(dim_in, dim_in))
for i in range(self.n_block):
setattr(self, 'ConvBlock'+str(i+1), ConvBlock(dim_in, dim_in))
setattr(self, 'ConvBlock' + str(i + 1), ConvBlock(dim_in, dim_in))
def forward(self, x):
skips = []
for i in range(self.n_skip):
skips.append(getattr(self, 'ConvBlockskip'+str(i+1))(x))
skips.append(getattr(self, 'ConvBlockskip' + str(i + 1))(x))
x = F.avg_pool2d(x, 2)
x = getattr(self, 'ConvBlock'+str(i+1))(x)
x = getattr(self, 'ConvBlock' + str(i + 1))(x)
x = self.ConvBlock5(x)
for i in range(self.n_skip):
x = getattr(self, 'ConvBlock'+str(i+6))(x)
x = getattr(self, 'ConvBlock' + str(i + 6))(x)
x = F.upsample(x, scale_factor=2)
x = skips[self.n_skip-i-1] + x
x = skips[self.n_skip - i - 1] + x
return x
......@@ -210,12 +199,9 @@ class HourGlass(nn.Layer):
self.use_res = use_res
self.HG = nn.Sequential(
HourGlassBlock(dim_in),
ConvBlock(dim_out, dim_out),
HourGlassBlock(dim_in), ConvBlock(dim_out, dim_out),
nn.Conv2D(dim_out, dim_out, kernel_size=1, stride=1, bias_attr=False),
nn.InstanceNorm2D(dim_out, weight_attr=False, bias_attr=False),
nn.ReLU()
)
nn.InstanceNorm2D(dim_out, weight_attr=False, bias_attr=False), nn.ReLU())
self.Conv1 = nn.Conv2D(dim_out, 3, kernel_size=1, stride=1)
......@@ -292,12 +278,12 @@ class SoftAdaLIN(nn.Layer):
self.w_gamma = self.create_parameter([1, num_features], default_initializer=nn.initializer.Constant(0.))
self.w_beta = self.create_parameter([1, num_features], default_initializer=nn.initializer.Constant(0.))
self.c_gamma = nn.Sequential(nn.Linear(num_features, num_features, bias_attr=False),
nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.c_beta = nn.Sequential(nn.Linear(num_features, num_features, bias_attr=False),
nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.c_gamma = nn.Sequential(
nn.Linear(num_features, num_features, bias_attr=False), nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.c_beta = nn.Sequential(
nn.Linear(num_features, num_features, bias_attr=False), nn.ReLU(),
nn.Linear(num_features, num_features, bias_attr=False))
self.s_gamma = nn.Linear(num_features, num_features, bias_attr=False)
self.s_beta = nn.Linear(num_features, num_features, bias_attr=False)
......@@ -352,8 +338,8 @@ class LIN(nn.Layer):
out = out * self.gamma.expand([x.shape[0], -1, -1, -1]) + self.beta.expand([x.shape[0], -1, -1, -1])
return out
if __name__ == '__main__':
#d = Discriminator(3)
# paddle.summary(d, (4, 3, 256, 256))
......
......@@ -8,13 +8,14 @@ import paddlehub as hub
from Photo2Cartoon.model import ResnetGenerator
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="Photo2Cartoon", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="Photo2Cartoon", # 模型介绍
version="1.0.0" # 版本号
name="Photo2Cartoon", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="Photo2Cartoon", # 模型介绍
version="1.0.0" # 版本号
)
class Photo2Cartoon(nn.Layer):
def __init__(self):
......@@ -49,18 +50,15 @@ class Photo2Cartoon(nn.Layer):
if images is not None:
datas = images
# 返回数据列表
return datas
# 数据预处理函数
def preprocess(self, images, batch_size, use_gpu):
# 获取人脸关键点
outputs = self.face_detector.keypoint_detection(
images=images,
batch_size=batch_size,
use_gpu=use_gpu)
outputs = self.face_detector.keypoint_detection(images=images, batch_size=batch_size, use_gpu=use_gpu)
crops = []
for output, image in zip(outputs, images):
for landmarks in output['data']:
......@@ -69,7 +67,8 @@ class Photo2Cartoon(nn.Layer):
# rotation angle
left_eye_corner = landmarks[36]
right_eye_corner = landmarks[45]
radian = np.arctan((left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
radian = np.arctan(
(left_eye_corner[1] - right_eye_corner[1]) / (left_eye_corner[0] - right_eye_corner[0]))
# image size after rotating
height, width, _ = image.shape
......@@ -84,7 +83,7 @@ class Photo2Cartoon(nn.Layer):
# affine matrix
M = np.array([[cos, sin, (1 - cos) * width / 2. - sin * height / 2. + Tx],
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
[-sin, cos, sin * width / 2. + (1 - cos) * height / 2. + Ty]])
image = cv2.warpAffine(image, M, (new_w, new_h), borderValue=(255, 255, 255))
......@@ -114,16 +113,17 @@ class Photo2Cartoon(nn.Layer):
h, w = image.shape[:2]
left_white = max(0, -left)
left = max(0, left)
right = min(right, w-1)
right_white = left_white + (right-left)
right = min(right, w - 1)
right_white = left_white + (right - left)
top_white = max(0, -top)
top = max(0, top)
bottom = min(bottom, h-1)
bottom = min(bottom, h - 1)
bottom_white = top_white + (bottom - top)
image_crop[top_white:bottom_white+1, left_white:right_white+1] = image[top:bottom+1, left:right+1].copy()
image_crop[top_white:bottom_white + 1, left_white:right_white + 1] = image[top:bottom + 1, left:right +
1].copy()
crops.append(image_crop)
# 获取人像分割的输出
results = self.seg.Segmentation(images=crops, batch_size=batch_size)
......@@ -158,7 +158,7 @@ class Photo2Cartoon(nn.Layer):
# 切分数据
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
return input_datas, masks
......@@ -175,7 +175,7 @@ class Photo2Cartoon(nn.Layer):
cartoon = self.net(data)
outputs.append(cartoon[0].numpy())
outputs = np.concatenate(outputs, 0)
return outputs
......@@ -204,17 +204,16 @@ class Photo2Cartoon(nn.Layer):
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), cartoon)
cartoons.append(cartoon)
return cartoons
def Cartoon_GEN(
self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
def Cartoon_GEN(self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False,
use_gpu=False):
# 获取输入数据
images = self.load_datas(paths, images)
......@@ -228,4 +227,4 @@ class Photo2Cartoon(nn.Layer):
# 结果后处理
cartoons = self.postprocess(outputs, masks, visualization, output_dir)
return cartoons
\ No newline at end of file
return cartoons
......@@ -6,18 +6,19 @@ from U2Net_Portrait.u2net import U2NET
from U2Net_Portrait.processor import Processor
from paddlehub.module.module import moduleinfo
@moduleinfo(
name="U2Net_Portrait", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Net_Portrait", # 模型介绍
version="1.0.0" # 版本号
name="U2Net_Portrait", # 模型名称
type="CV", # 模型类型
author="jm12138", # 作者名称
author_email="jm12138@qq.com", # 作者邮箱
summary="U2Net_Portrait", # 模型介绍
version="1.0.0" # 版本号
)
class U2Net_Portrait(nn.Layer):
def __init__(self):
super(U2Net_Portrait, self).__init__()
self.model = U2NET(3,1)
self.model = U2NET(3, 1)
state_dict = paddle.load(os.path.join(self.directory, 'u2net_portrait.pdparams'))
self.model.set_dict(state_dict)
self.model.eval()
......@@ -26,22 +27,21 @@ class U2Net_Portrait(nn.Layer):
outputs = []
for data in input_datas:
data = paddle.to_tensor(data, 'float32')
d1,d2,d3,d4,d5,d6,d7= self.model(data)
d1, d2, d3, d4, d5, d6, d7 = self.model(data)
outputs.append(d1.numpy())
outputs = np.concatenate(outputs, 0)
return outputs
def Portrait_GEN(
self,
images=None,
paths=None,
scale=1,
batch_size=1,
output_dir='output',
face_detection=True,
visualization=False):
def Portrait_GEN(self,
images=None,
paths=None,
scale=1,
batch_size=1,
output_dir='output',
face_detection=True,
visualization=False):
# 初始化数据处理器
processor = Processor(paths, images, batch_size, face_detection, scale)
......
......@@ -5,6 +5,7 @@ import paddlehub as hub
__all__ = ['Processor']
class Processor():
def __init__(self, paths, images, batch_size, face_detection=True, scale=1):
# 图像列表
......@@ -26,7 +27,7 @@ class Processor():
if images is not None:
datas = images
# 返回数据列表
return datas
......@@ -35,10 +36,7 @@ class Processor():
if face_detection:
# face detection
face_detector = hub.Module(name="pyramidbox_lite_mobile")
results = face_detector.face_detection(images=imgs,
use_gpu=False,
visualization=False,
confs_threshold=0.5)
results = face_detector.face_detection(images=imgs, use_gpu=False, visualization=False, confs_threshold=0.5)
im_faces = []
for datas, img in zip(results, imgs):
for face in datas['data']:
......@@ -46,52 +44,64 @@ class Processor():
l, r, t, b = [face['left'], face['right'], face['top'], face['bottom']]
# square crop
pad = max(int(scale*(r-l)), int(scale*(b-t)))
c_w, c_h = (r-l)//2+l, (b-t)//2+t
top = 0 if c_h-pad<0 else c_h-pad
pad = max(int(scale * (r - l)), int(scale * (b - t)))
c_w, c_h = (r - l) // 2 + l, (b - t) // 2 + t
top = 0 if c_h - pad < 0 else c_h - pad
bottom = pad + c_h
left = 0 if c_w-pad<0 else c_w-pad
left = 0 if c_w - pad < 0 else c_w - pad
right = pad + c_w
crop = img[top:bottom, left:right]
# resize
im_face = cv2.resize(crop, (512,512), interpolation = cv2.INTER_AREA)
im_face = cv2.resize(crop, (512, 512), interpolation=cv2.INTER_AREA)
im_faces.append(im_face)
else:
im_faces = []
for img in imgs:
h, w = img.shape[:2]
if h>w:
if (h-w)%2==0:
img = np.pad(img,((0,0),((h-w)//2,(h-w)//2),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
if h > w:
if (h - w) % 2 == 0:
img = np.pad(
img, ((0, 0), ((h - w) // 2, (h - w) // 2), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
else:
img = np.pad(img,((0,0),((h-w)//2,(h-w)//2+1),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
img = np.pad(
img, ((0, 0), ((h - w) // 2, (h - w) // 2 + 1), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
else:
if (w-h)%2==0:
img = np.pad(img,(((w-h)//2,(w-h)//2),(0,0),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
if (w - h) % 2 == 0:
img = np.pad(
img, (((w - h) // 2, (w - h) // 2), (0, 0), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
else:
img = np.pad(img,(((w-h)//2,(w-h)//2+1),(0,0),(0,0)),mode='constant',constant_values=((255,255),(255,255),(255,255)))
im_face = cv2.resize(img, (512,512), interpolation = cv2.INTER_AREA)
img = np.pad(
img, (((w - h) // 2, (w - h) // 2 + 1), (0, 0), (0, 0)),
mode='constant',
constant_values=((255, 255), (255, 255), (255, 255)))
im_face = cv2.resize(img, (512, 512), interpolation=cv2.INTER_AREA)
im_faces.append(im_face)
input_datas = []
for im_face in im_faces:
tmpImg = np.zeros((im_face.shape[0],im_face.shape[1],3))
im_face = im_face/np.max(im_face)
tmpImg = np.zeros((im_face.shape[0], im_face.shape[1], 3))
im_face = im_face / np.max(im_face)
tmpImg[:,:,0] = (im_face[:,:,2]-0.406)/0.225
tmpImg[:,:,1] = (im_face[:,:,1]-0.456)/0.224
tmpImg[:,:,2] = (im_face[:,:,0]-0.485)/0.229
tmpImg[:, :, 0] = (im_face[:, :, 2] - 0.406) / 0.225
tmpImg[:, :, 1] = (im_face[:, :, 1] - 0.456) / 0.224
tmpImg[:, :, 2] = (im_face[:, :, 0] - 0.485) / 0.229
# convert BGR to RGB
tmpImg = tmpImg.transpose((2, 0, 1))
tmpImg = tmpImg[np.newaxis,:,:,:]
tmpImg = tmpImg[np.newaxis, :, :, :]
input_datas.append(tmpImg)
input_datas = np.concatenate(input_datas, 0)
datas_num = input_datas.shape[0]
split_num = datas_num//batch_size+1 if datas_num%batch_size!=0 else datas_num//batch_size
split_num = datas_num // batch_size + 1 if datas_num % batch_size != 0 else datas_num // batch_size
input_datas = np.array_split(input_datas, split_num)
......@@ -101,7 +111,7 @@ class Processor():
ma = np.max(d)
mi = np.min(d)
dn = (d-mi)/(ma-mi)
dn = (d - mi) / (ma - mi)
return dn
......@@ -110,20 +120,20 @@ class Processor():
results = []
if visualization and not os.path.exists(output_dir):
os.mkdir(output_dir)
for i in range(outputs.shape[0]):
# normalization
pred = 1.0 - outputs[i,0,:,:]
pred = 1.0 - outputs[i, 0, :, :]
pred = self.normPRED(pred)
# convert torch tensor to numpy array
pred = pred.squeeze()
pred = (pred*255).astype(np.uint8)
pred = (pred * 255).astype(np.uint8)
results.append(pred)
if visualization:
cv2.imwrite(os.path.join(output_dir, 'result_%d.png' % i), pred)
return results
\ No newline at end of file
return results
import paddle
import paddle.nn as nn
import paddle.nn.functional as F
__all__ = ['U2NETP', 'U2NET']
class REBNCONV(nn.Layer):
def __init__(self, in_ch=3, out_ch=3, dirate=1):
super(REBNCONV, self).__init__()
self.conv_s1 = nn.Conv2D(in_ch, out_ch, 3, padding=1 * dirate, dilation=1 * dirate)
self.bn_s1 = nn.BatchNorm2D(out_ch)
self.relu_s1 = nn.ReLU()
def forward(self, x):
hx = x
xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
return xout
## upsample tensor 'src' to have the same spatial size with tensor 'tar'
def _upsample_like(src, tar):
src = F.upsample(src, size=tar.shape[2:], mode='bilinear')
return src
### RSU-7 ###
class RSU7(nn.Layer): #UNet07DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU7, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool5 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv7 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv6d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx = self.pool5(hx5)
hx6 = self.rebnconv6(hx)
hx7 = self.rebnconv7(hx6)
hx6d = self.rebnconv6d(paddle.concat((hx7, hx6), 1))
hx6dup = _upsample_like(hx6d, hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6dup, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-6 ###
class RSU6(nn.Layer): #UNet06DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU6, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool4 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv6 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv5d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx = self.pool4(hx4)
hx5 = self.rebnconv5(hx)
hx6 = self.rebnconv6(hx5)
hx5d = self.rebnconv5d(paddle.concat((hx6, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-5 ###
class RSU5(nn.Layer): #UNet05DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU5, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool3 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv5 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv4d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx = self.pool3(hx3)
hx4 = self.rebnconv4(hx)
hx5 = self.rebnconv5(hx4)
hx4d = self.rebnconv4d(paddle.concat((hx5, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4 ###
class RSU4(nn.Layer): #UNet04DRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.pool1 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.pool2 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=1)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=1)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx = self.pool1(hx1)
hx2 = self.rebnconv2(hx)
hx = self.pool2(hx2)
hx3 = self.rebnconv3(hx)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.rebnconv2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.rebnconv1d(paddle.concat((hx2dup, hx1), 1))
return hx1d + hxin
### RSU-4F ###
class RSU4F(nn.Layer): #UNet04FRES(nn.Layer):
def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
super(RSU4F, self).__init__()
self.rebnconvin = REBNCONV(in_ch, out_ch, dirate=1)
self.rebnconv1 = REBNCONV(out_ch, mid_ch, dirate=1)
self.rebnconv2 = REBNCONV(mid_ch, mid_ch, dirate=2)
self.rebnconv3 = REBNCONV(mid_ch, mid_ch, dirate=4)
self.rebnconv4 = REBNCONV(mid_ch, mid_ch, dirate=8)
self.rebnconv3d = REBNCONV(mid_ch * 2, mid_ch, dirate=4)
self.rebnconv2d = REBNCONV(mid_ch * 2, mid_ch, dirate=2)
self.rebnconv1d = REBNCONV(mid_ch * 2, out_ch, dirate=1)
def forward(self, x):
hx = x
hxin = self.rebnconvin(hx)
hx1 = self.rebnconv1(hxin)
hx2 = self.rebnconv2(hx1)
hx3 = self.rebnconv3(hx2)
hx4 = self.rebnconv4(hx3)
hx3d = self.rebnconv3d(paddle.concat((hx4, hx3), 1))
hx2d = self.rebnconv2d(paddle.concat((hx3d, hx2), 1))
hx1d = self.rebnconv1d(paddle.concat((hx2d, hx1), 1))
return hx1d + hxin
##### U^2-Net ####
class U2NET(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NET, self).__init__()
self.stage1 = RSU7(in_ch, 32, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 32, 128)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(128, 64, 256)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(256, 128, 512)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(512, 256, 512)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(512, 256, 512)
# decoder
self.stage5d = RSU4F(1024, 256, 512)
self.stage4d = RSU4(1024, 128, 256)
self.stage3d = RSU5(512, 64, 128)
self.stage2d = RSU6(256, 32, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(128, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(256, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(512, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(512, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#-------------------- decoder --------------------
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
### U^2-Net small ###
class U2NETP(nn.Layer):
def __init__(self, in_ch=3, out_ch=1):
super(U2NETP, self).__init__()
self.stage1 = RSU7(in_ch, 16, 64)
self.pool12 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage2 = RSU6(64, 16, 64)
self.pool23 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage3 = RSU5(64, 16, 64)
self.pool34 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage4 = RSU4(64, 16, 64)
self.pool45 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage5 = RSU4F(64, 16, 64)
self.pool56 = nn.MaxPool2D(2, stride=2, ceil_mode=True)
self.stage6 = RSU4F(64, 16, 64)
# decoder
self.stage5d = RSU4F(128, 16, 64)
self.stage4d = RSU4(128, 16, 64)
self.stage3d = RSU5(128, 16, 64)
self.stage2d = RSU6(128, 16, 64)
self.stage1d = RSU7(128, 16, 64)
self.side1 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side2 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side3 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side4 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side5 = nn.Conv2D(64, out_ch, 3, padding=1)
self.side6 = nn.Conv2D(64, out_ch, 3, padding=1)
self.outconv = nn.Conv2D(6, out_ch, 1)
def forward(self, x):
hx = x
#stage 1
hx1 = self.stage1(hx)
hx = self.pool12(hx1)
#stage 2
hx2 = self.stage2(hx)
hx = self.pool23(hx2)
#stage 3
hx3 = self.stage3(hx)
hx = self.pool34(hx3)
#stage 4
hx4 = self.stage4(hx)
hx = self.pool45(hx4)
#stage 5
hx5 = self.stage5(hx)
hx = self.pool56(hx5)
#stage 6
hx6 = self.stage6(hx)
hx6up = _upsample_like(hx6, hx5)
#decoder
hx5d = self.stage5d(paddle.concat((hx6up, hx5), 1))
hx5dup = _upsample_like(hx5d, hx4)
hx4d = self.stage4d(paddle.concat((hx5dup, hx4), 1))
hx4dup = _upsample_like(hx4d, hx3)
hx3d = self.stage3d(paddle.concat((hx4dup, hx3), 1))
hx3dup = _upsample_like(hx3d, hx2)
hx2d = self.stage2d(paddle.concat((hx3dup, hx2), 1))
hx2dup = _upsample_like(hx2d, hx1)
hx1d = self.stage1d(paddle.concat((hx2dup, hx1), 1))
#side output
d1 = self.side1(hx1d)
d2 = self.side2(hx2d)
d2 = _upsample_like(d2, d1)
d3 = self.side3(hx3d)
d3 = _upsample_like(d3, d1)
d4 = self.side4(hx4d)
d4 = _upsample_like(d4, d1)
d5 = self.side5(hx5d)
d5 = _upsample_like(d5, d1)
d6 = self.side6(hx6)
d6 = _upsample_like(d6, d1)
d0 = self.outconv(paddle.concat((d1, d2, d3, d4, d5, d6), 1))
return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
......@@ -24,9 +25,11 @@ class Model():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -36,7 +39,7 @@ class Model():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -47,10 +50,10 @@ class Model():
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
......@@ -65,9 +68,9 @@ class Model():
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['InferenceModel']
class InferenceModel():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True):
......@@ -25,14 +26,10 @@ class InferenceModel():
# 打印函数
def __repr__(self):
'''
get the numbers and name of inputs and outputs
get the numbers and name of inputs and outputs
'''
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (
len(self.input_handles),
str(self.input_names),
len(self.output_handles),
str(self.output_names)
)
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (len(
self.input_handles), str(self.input_names), len(self.output_handles), str(self.output_names))
# 类调用函数
def __call__(self, *input_datas, batch_size=1):
......@@ -59,9 +56,11 @@ class InferenceModel():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -71,7 +70,7 @@ class InferenceModel():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -125,17 +124,17 @@ class InferenceModel():
self.input_handles[i].copy_from_cpu(input_data)
self.predictor.run()
for i in range(len(self.output_handles)):
output = self.output_handles[i].copy_to_cpu()
if i in outputs:
outputs[i].append(output)
else:
outputs[i] = [output]
# 预测结果合并
for key in outputs.keys():
outputs[key] = np.concatenate(outputs[key], 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -26,12 +26,7 @@ class MiDaS_Large(Module):
model_path = os.path.join(self.directory, "model-f6b98070")
# 加载模型
self.model = InferenceModel(
modelpath=model_path,
use_gpu=use_gpu,
use_mkldnn=False,
combined=True
)
self.model = InferenceModel(modelpath=model_path, use_gpu=use_gpu, use_mkldnn=False, combined=True)
self.model.eval()
# 数据预处理配置
......@@ -46,11 +41,10 @@ class MiDaS_Large(Module):
resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC,
),
NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
PrepareForNet()
])
# 数据读取函数
@staticmethod
def load_datas(paths, images):
......@@ -84,7 +78,7 @@ class MiDaS_Large(Module):
input_data = img[np.newaxis, ...]
input_datas.append(input_data)
# 拼接数据
input_datas = np.concatenate(input_datas, 0)
......@@ -113,12 +107,7 @@ class MiDaS_Large(Module):
return outputs
# 深度估计函数
def depth_estimation(self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
def depth_estimation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
# 加载数据
datas = self.load_datas(paths, images)
......@@ -130,5 +119,5 @@ class MiDaS_Large(Module):
# 结果后处理
outputs = self.postprocess(datas, results, output_dir, visualization)
return outputs
\ No newline at end of file
return outputs
......@@ -7,6 +7,7 @@ import cv2
class Resize(object):
"""Resize sample to given size (width, height).
"""
def __init__(self,
width,
height,
......@@ -51,12 +52,10 @@ class Resize(object):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
if y < min_val:
y = (np.ceil(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
return y
......@@ -91,31 +90,24 @@ class Resize(object):
# fit height
scale_width = scale_height
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
if self.__resize_method == "lower_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
elif self.__resize_method == "upper_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
elif self.__resize_method == "minimal":
new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width)
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
return (new_width, new_height)
def __call__(self, sample):
width, height = self.get_size(sample["image"].shape[1],
sample["image"].shape[0])
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
# resize sample
sample["image"] = cv2.resize(
......@@ -133,8 +125,7 @@ class Resize(object):
)
if "depth" in sample:
sample["depth"] = cv2.resize(sample["depth"], (width, height),
interpolation=cv2.INTER_NEAREST)
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
sample["mask"] = cv2.resize(
sample["mask"].astype(np.float32),
......@@ -149,6 +140,7 @@ class Resize(object):
class NormalizeImage(object):
"""Normlize image by given mean and std.
"""
def __init__(self, mean, std):
self.__mean = mean
self.__std = std
......@@ -162,6 +154,7 @@ class NormalizeImage(object):
class PrepareForNet(object):
"""Prepare sample for usage as network input.
"""
def __init__(self):
pass
......
......@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (len(image.shape) == 2
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
elif (len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False
else:
raise Exception(
"Image must have H x W x 3, H x W x 1 or H x W dimensions.")
raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
file.write("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
......
......@@ -5,6 +5,7 @@ from paddle.inference import create_predictor, Config
__all__ = ['InferenceModel']
class InferenceModel():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=False, combined=True):
......@@ -25,14 +26,10 @@ class InferenceModel():
# 打印函数
def __repr__(self):
'''
get the numbers and name of inputs and outputs
get the numbers and name of inputs and outputs
'''
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (
len(self.input_handles),
str(self.input_names),
len(self.output_handles),
str(self.output_names)
)
return 'inputs_num: %d\ninputs_names: %s\noutputs_num: %d\noutputs_names: %s' % (len(
self.input_handles), str(self.input_names), len(self.output_handles), str(self.output_names))
# 类调用函数
def __call__(self, *input_datas, batch_size=1):
......@@ -59,9 +56,11 @@ class InferenceModel():
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print('Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.')
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
......@@ -71,7 +70,7 @@ class InferenceModel():
config = Config(modelpath)
# 设置参数
if use_gpu:
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
......@@ -125,17 +124,17 @@ class InferenceModel():
self.input_handles[i].copy_from_cpu(input_data)
self.predictor.run()
for i in range(len(self.output_handles)):
output = self.output_handles[i].copy_to_cpu()
if i in outputs:
outputs[i].append(output)
else:
outputs[i] = [output]
# 预测结果合并
for key in outputs.keys():
outputs[key] = np.concatenate(outputs[key], 0)
# 返回预测结果
return outputs
\ No newline at end of file
return outputs
......@@ -26,12 +26,7 @@ class MiDaS_Small(Module):
model_path = os.path.join(self.directory, "model-small")
# 加载模型
self.model = InferenceModel(
modelpath=model_path,
use_gpu=use_gpu,
use_mkldnn=False,
combined=True
)
self.model = InferenceModel(modelpath=model_path, use_gpu=use_gpu, use_mkldnn=False, combined=True)
self.model.eval()
# 数据预处理配置
......@@ -46,11 +41,10 @@ class MiDaS_Small(Module):
resize_method="upper_bound",
image_interpolation_method=cv2.INTER_CUBIC,
),
NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
PrepareForNet()
])
# 数据读取函数
@staticmethod
def load_datas(paths, images):
......@@ -84,7 +78,7 @@ class MiDaS_Small(Module):
input_data = img[np.newaxis, ...]
input_datas.append(input_data)
# 拼接数据
input_datas = np.concatenate(input_datas, 0)
......@@ -113,12 +107,7 @@ class MiDaS_Small(Module):
return outputs
# 深度估计函数
def depth_estimation(self,
images=None,
paths=None,
batch_size=1,
output_dir='output',
visualization=False):
def depth_estimation(self, images=None, paths=None, batch_size=1, output_dir='output', visualization=False):
# 加载数据
datas = self.load_datas(paths, images)
......@@ -130,5 +119,5 @@ class MiDaS_Small(Module):
# 结果后处理
outputs = self.postprocess(datas, results, output_dir, visualization)
return outputs
\ No newline at end of file
return outputs
......@@ -7,6 +7,7 @@ import cv2
class Resize(object):
"""Resize sample to given size (width, height).
"""
def __init__(self,
width,
height,
......@@ -51,12 +52,10 @@ class Resize(object):
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
if max_val is not None and y > max_val:
y = (np.floor(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
if y < min_val:
y = (np.ceil(x / self.__multiple_of) *
self.__multiple_of).astype(int)
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
return y
......@@ -91,31 +90,24 @@ class Resize(object):
# fit height
scale_width = scale_height
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
if self.__resize_method == "lower_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
min_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, min_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, min_val=self.__width)
elif self.__resize_method == "upper_bound":
new_height = self.constrain_to_multiple_of(scale_height * height,
max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width,
max_val=self.__width)
new_height = self.constrain_to_multiple_of(scale_height * height, max_val=self.__height)
new_width = self.constrain_to_multiple_of(scale_width * width, max_val=self.__width)
elif self.__resize_method == "minimal":
new_height = self.constrain_to_multiple_of(scale_height * height)
new_width = self.constrain_to_multiple_of(scale_width * width)
else:
raise ValueError(
f"resize_method {self.__resize_method} not implemented")
raise ValueError(f"resize_method {self.__resize_method} not implemented")
return (new_width, new_height)
def __call__(self, sample):
width, height = self.get_size(sample["image"].shape[1],
sample["image"].shape[0])
width, height = self.get_size(sample["image"].shape[1], sample["image"].shape[0])
# resize sample
sample["image"] = cv2.resize(
......@@ -133,8 +125,7 @@ class Resize(object):
)
if "depth" in sample:
sample["depth"] = cv2.resize(sample["depth"], (width, height),
interpolation=cv2.INTER_NEAREST)
sample["depth"] = cv2.resize(sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST)
sample["mask"] = cv2.resize(
sample["mask"].astype(np.float32),
......@@ -149,6 +140,7 @@ class Resize(object):
class NormalizeImage(object):
"""Normlize image by given mean and std.
"""
def __init__(self, mean, std):
self.__mean = mean
self.__std = std
......@@ -162,6 +154,7 @@ class NormalizeImage(object):
class PrepareForNet(object):
"""Prepare sample for usage as network input.
"""
def __init__(self):
pass
......
......@@ -25,12 +25,10 @@ def write_pfm(path, image, scale=1):
if len(image.shape) == 3 and image.shape[2] == 3: # color image
color = True
elif (len(image.shape) == 2
or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
elif (len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1): # greyscale
color = False
else:
raise Exception(
"Image must have H x W x 3, H x W x 1 or H x W dimensions.")
raise Exception("Image must have H x W x 3, H x W x 1 or H x W dimensions.")
file.write("PF\n" if color else "Pf\n".encode())
file.write("%d %d\n".encode() % (image.shape[1], image.shape[0]))
......
import os
import numpy as np
from paddle.inference import create_predictor, Config
__all__ = ['Model']
class Model():
# 初始化函数
def __init__(self, modelpath, use_gpu=False, use_mkldnn=True, combined=True):
# 加载模型预测器
self.predictor = self.load_model(modelpath, use_gpu, use_mkldnn, combined)
# 获取模型的输入输出
self.input_names = self.predictor.get_input_names()
self.output_names = self.predictor.get_output_names()
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
# 模型加载函数
def load_model(self, modelpath, use_gpu, use_mkldnn, combined):
# 对运行位置进行配置
if use_gpu:
try:
int(os.environ.get('CUDA_VISIBLE_DEVICES'))
except Exception:
print(
'Error! Unable to use GPU. Please set the environment variables "CUDA_VISIBLE_DEVICES=GPU_id" to use GPU.'
)
use_gpu = False
# 加载模型参数
if combined:
model = os.path.join(modelpath, "__model__")
params = os.path.join(modelpath, "__params__")
config = Config(model, params)
else:
config = Config(modelpath)
# 设置参数
if use_gpu:
config.enable_use_gpu(100, 0)
else:
config.disable_gpu()
if use_mkldnn:
config.enable_mkldnn()
config.disable_glog_info()
config.switch_ir_optim(True)
config.enable_memory_optim()
config.switch_use_feed_fetch_ops(False)
config.switch_specify_input_names(True)
# 通过参数加载模型预测器
predictor = create_predictor(config)
# 返回预测器
return predictor
# 模型预测函数
def predict(self, input_datas):
outputs = []
# 遍历输入数据进行预测
for input_data in input_datas:
inputs = input_data.copy()
self.input_handle.copy_from_cpu(inputs)
self.predictor.run()
output = self.output_handle.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
# 返回预测结果
return outputs
......@@ -2,10 +2,11 @@ import numpy as np
import cv2
from scipy import ndimage
def get_normal_map(img):
img = img.astype(np.float)
img = img / 255.0
img = - img + 1
img = -img + 1
img[img < 0] = 0
img[img > 1] = 1
return img
......@@ -112,7 +113,7 @@ def resize_img_512(img):
def resize_img_512_3d(img):
zeros = np.zeros((1, 3, 512, 512), dtype=np.float)
zeros[0, 0: img.shape[0], 0: img.shape[1], 0: img.shape[2]] = img
zeros[0, 0:img.shape[0], 0:img.shape[1], 0:img.shape[2]] = img
return zeros.transpose((1, 2, 3, 0))
......@@ -122,7 +123,7 @@ def denoise_mat(img, i):
def show_active_img_and_save_denoise(img, path):
mat = img.astype(np.float)
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......@@ -132,11 +133,9 @@ def show_active_img_and_save_denoise(img, path):
return
def show_active_img(name, img):
mat = img.astype(np.float)
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......@@ -147,7 +146,7 @@ def show_active_img(name, img):
def get_active_img(img):
mat = img.astype(np.float)
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......@@ -158,7 +157,7 @@ def get_active_img(img):
def get_active_img_fil(img):
mat = img.astype(np.float)
mat[mat < 0.18] = 0
mat = - mat + 1
mat = -mat + 1
mat = mat * 255.0
mat[mat < 0] = 0
mat[mat > 255] = 255
......
......@@ -14,12 +14,12 @@ import paddle.fluid as fluid
import paddlehub as hub
from Extract_Line_Draft.function import *
@moduleinfo(
name="Extract_Line_Draft",
version="1.0.0",
type="cv/segmentation",
summary=
"Import the color picture and generate the line draft of the picture",
summary="Import the color picture and generate the line draft of the picture",
author="彭兆帅,郑博培",
author_email="1084667371@qq.com,2733821739@qq.com")
class ExtractLineDraft(hub.Module):
......@@ -28,9 +28,10 @@ class ExtractLineDraft(hub.Module):
Initialize with the necessary elements
"""
# 加载模型路径
self.default_pretrained_model_path = os.path.join(self.directory, "assets","infer_model")
self.default_pretrained_model_path = os.path.join(self.directory, "assets", "infer_model")
self._set_config()
def _set_config(self):
def _set_config(self):
"""
predictor config setting
"""
......@@ -72,7 +73,7 @@ class ExtractLineDraft(hub.Module):
self.predictor.zero_copy_run()
output = self.output_tensor.copy_to_cpu()
outputs.append(output)
# 预测结果合并
outputs = np.concatenate(outputs, 0)
......@@ -153,12 +154,9 @@ class ExtractLineDraft(hub.Module):
usage='%(prog)s',
add_help=True)
self.arg_input_group = self.parser.add_argument_group(
title="Input options", description="Input data. Required")
self.arg_input_group = self.parser.add_argument_group(title="Input options", description="Input data. Required")
self.arg_config_group = self.parser.add_argument_group(
title="Config options",
description=
"Run configuration for controlling module behavior, not required.")
title="Config options", description="Run configuration for controlling module behavior, not required.")
self.add_module_input_arg()
......@@ -173,21 +171,12 @@ class ExtractLineDraft(hub.Module):
use_gpu = args.use_gpu
self.ExtractLine(image=input_data, use_gpu=use_gpu)
def add_module_input_arg(self):
"""
Add the command input options
"""
self.arg_input_group.add_argument(
'--image',
type=str,
default=None,
help="file contain input data")
self.arg_input_group.add_argument(
'--use_gpu',
type=ast.literal_eval,
default=None,
help="weather to use gpu")
self.arg_input_group.add_argument('--image', type=str, default=None, help="file contain input data")
self.arg_input_group.add_argument('--use_gpu', type=ast.literal_eval, default=None, help="weather to use gpu")
def check_input_data(self, args):
input_data = []
......@@ -196,4 +185,3 @@ class ExtractLineDraft(hub.Module):
raise RuntimeError("Path %s is not exist." % args.image)
path = "{}".format(args.image)
return path
......@@ -48,12 +48,9 @@ class FCN(nn.Layer):
super(FCN, self).__init__()
self.backbone = backbone
backbone_channels = [
backbone.feat_channels[i] for i in backbone_indices
]
backbone_channels = [backbone.feat_channels[i] for i in backbone_indices]
self.head = FCNHead(num_classes, backbone_indices, backbone_channels,
channels)
self.head = FCNHead(num_classes, backbone_indices, backbone_channels, channels)
self.align_corners = align_corners
self.pretrained = pretrained
......@@ -62,11 +59,7 @@ class FCN(nn.Layer):
feat_list = self.backbone(x)
logit_list = self.head(feat_list)
return [
F.interpolate(
logit,
x.shape[2:],
mode='bilinear',
align_corners=self.align_corners) for logit in logit_list
F.interpolate(logit, x.shape[2:], mode='bilinear', align_corners=self.align_corners) for logit in logit_list
]
......@@ -83,11 +76,7 @@ class FCNHead(nn.Layer):
pretrained (str, optional): The path of pretrained model. Default: None
"""
def __init__(self,
num_classes,
backbone_indices=(-1, ),
backbone_channels=(270, ),
channels=None):
def __init__(self, num_classes, backbone_indices=(-1, ), backbone_channels=(270, ), channels=None):
super(FCNHead, self).__init__()
self.num_classes = num_classes
......@@ -96,17 +85,8 @@ class FCNHead(nn.Layer):
channels = backbone_channels[0]
self.conv_1 = ConvBNReLU(
in_channels=backbone_channels[0],
out_channels=channels,
kernel_size=1,
padding='same',
stride=1)
self.cls = nn.Conv2D(
in_channels=channels,
out_channels=self.num_classes,
kernel_size=1,
stride=1,
padding=0)
in_channels=backbone_channels[0], out_channels=channels, kernel_size=1, padding='same', stride=1)
self.cls = nn.Conv2D(in_channels=channels, out_channels=self.num_classes, kernel_size=1, stride=1, padding=0)
def forward(self, feat_list):
logit_list = []
......@@ -115,4 +95,3 @@ class FCNHead(nn.Layer):
logit = self.cls(x)
logit_list.append(logit)
return logit_list
......@@ -26,16 +26,10 @@ def SyncBatchNorm(*args, **kwargs):
class ConvBNReLU(nn.Layer):
def __init__(self,
in_channels,
out_channels,
kernel_size,
padding='same',
**kwargs):
def __init__(self, in_channels, out_channels, kernel_size, padding='same', **kwargs):
super().__init__()
self._conv = nn.Conv2D(
in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._batch_norm = SyncBatchNorm(out_channels)
......@@ -47,19 +41,12 @@ class ConvBNReLU(nn.Layer):
class ConvBN(nn.Layer):
def __init__(self,
in_channels,
out_channels,
kernel_size,
padding='same',
**kwargs):
def __init__(self, in_channels, out_channels, kernel_size, padding='same', **kwargs):
super().__init__()
self._conv = nn.Conv2D(
in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._conv = nn.Conv2D(in_channels, out_channels, kernel_size, padding=padding, **kwargs)
self._batch_norm = SyncBatchNorm(out_channels)
def forward(self, x):
x = self._conv(x)
x = self._batch_norm(x)
return x
此差异已折叠。
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册