Skip to content

P
PaddleHub

PaddlePaddle / PaddleHub
1 年多 前同步成功

通知 283
Star 12117
Fork 2091
P
PaddleHub
未验证 提交 821d6c14 编写于 作者: H haoyuying 提交者： GitHub
浏览文件
操作

2.0rc for msgnet and user_guided_colorization

上级 28bde845
隐藏空白更改
内联 并排
Showing with 97 addition and 103 deletion
hub_module/modules/image/colorization/user_guided_colorization/module.py
浏览文件 @ 821d6c14
......@@ -16,7 +16,7 @@ import os
import paddle
import paddle.nn as nn
from paddle.nn import Conv2d, ConvTranspose2d
from paddle.nn import Conv2D, Conv2DTranspose
from paddlehub.module.module import moduleinfo
import paddlehub.process.transforms as T
from paddlehub.module.cv_module import ImageColorizeModule
......@@ -48,110 +48,110 @@ class UserGuidedColorization(nn.Layer):
self.classification = classification
# Conv1
model1 = (
Conv2d(self.input_nc, 64, 3, 1, 1),
Conv2D(self.input_nc, 64, 3, 1, 1),
nn.ReLU(),
Conv2d(64, 64, 3, 1, 1),
Conv2D(64, 64, 3, 1, 1),
nn.ReLU(),
nn.BatchNorm(64),
)
# Conv2
model2 = (
Conv2d(64, 128, 3, 1, 1),
Conv2D(64, 128, 3, 1, 1),
nn.ReLU(),
Conv2d(128, 128, 3, 1, 1),
Conv2D(128, 128, 3, 1, 1),
nn.ReLU(),
nn.BatchNorm(128),
)
# Conv3
model3 = (
Conv2d(128, 256, 3, 1, 1),
Conv2D(128, 256, 3, 1, 1),
nn.ReLU(),
Conv2d(256, 256, 3, 1, 1),
Conv2D(256, 256, 3, 1, 1),
nn.ReLU(),
Conv2d(256, 256, 3, 1, 1),
Conv2D(256, 256, 3, 1, 1),
nn.ReLU(),
nn.BatchNorm(256),
)
# Conv4
model4 = (
Conv2d(256, 512, 3, 1, 1),
Conv2D(256, 512, 3, 1, 1),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 1),
Conv2D(512, 512, 3, 1, 1),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 1),
Conv2D(512, 512, 3, 1, 1),
nn.ReLU(),
nn.BatchNorm(512),
)
# Conv5
model5 = (
Conv2d(512, 512, 3, 1, 2, 2),
Conv2D(512, 512, 3, 1, 2, 2),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 2, 2),
Conv2D(512, 512, 3, 1, 2, 2),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 2, 2),
Conv2D(512, 512, 3, 1, 2, 2),
nn.ReLU(),
nn.BatchNorm(512),
)
# Conv6
model6 = (
Conv2d(512, 512, 3, 1, 2, 2),
Conv2D(512, 512, 3, 1, 2, 2),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 2, 2),
Conv2D(512, 512, 3, 1, 2, 2),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 2, 2),
Conv2D(512, 512, 3, 1, 2, 2),
nn.ReLU(),
nn.BatchNorm(512),
)
# Conv7
model7 = (
Conv2d(512, 512, 3, 1, 1),
Conv2D(512, 512, 3, 1, 1),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 1),
Conv2D(512, 512, 3, 1, 1),
nn.ReLU(),
Conv2d(512, 512, 3, 1, 1),
Conv2D(512, 512, 3, 1, 1),
nn.ReLU(),
nn.BatchNorm(512),
)
# Conv8
model8up = (ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1), )
model3short8 = (Conv2d(256, 256, 3, 1, 1), )
model8up = (Conv2DTranspose(512, 256, kernel_size=4, stride=2, padding=1), )
model3short8 = (Conv2D(256, 256, 3, 1, 1), )
model8 = (
nn.ReLU(),
Conv2d(256, 256, 3, 1, 1),
Conv2D(256, 256, 3, 1, 1),
nn.ReLU(),
Conv2d(256, 256, 3, 1, 1),
Conv2D(256, 256, 3, 1, 1),
nn.ReLU(),
nn.BatchNorm(256),
)
# Conv9
model9up = (ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1), )
model2short9 = (Conv2d(
model9up = (Conv2DTranspose(256, 128, kernel_size=4, stride=2, padding=1), )
model2short9 = (Conv2D(
128,
128,
3,
1,
1,
), )
model9 = (nn.ReLU(), Conv2d(128, 128, 3, 1, 1), nn.ReLU(), nn.BatchNorm(128))
model9 = (nn.ReLU(), Conv2D(128, 128, 3, 1, 1), nn.ReLU(), nn.BatchNorm(128))
# Conv10
model10up = (ConvTranspose2d(128, 128, kernel_size=4, stride=2, padding=1), )
model1short10 = (Conv2d(64, 128, 3, 1, 1), )
model10 = (nn.ReLU(), Conv2d(128, 128, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2))
model_class = (Conv2d(256, 529, 1), )
model10up = (Conv2DTranspose(128, 128, kernel_size=4, stride=2, padding=1), )
model1short10 = (Conv2D(64, 128, 3, 1, 1), )
model10 = (nn.ReLU(), Conv2D(128, 128, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2))
model_class = (Conv2D(256, 529, 1), )
if use_tanh:
model_out = (Conv2d(128, 2, 1, 1, 0, 1), nn.Tanh())
model_out = (Conv2D(128, 2, 1, 1, 0, 1), nn.Tanh())
else:
model_out = (Conv2d(128, 2, 1, 1, 0, 1), )
model_out = (Conv2D(128, 2, 1, 1, 0, 1), )
self.model1 = nn.Sequential(*model1)
self.model2 = nn.Sequential(*model2)
......@@ -173,7 +173,7 @@ class UserGuidedColorization(nn.Layer):
self.model_out = nn.Sequential(*model_out)
if load_checkpoint is not None:
model_dict = paddle.load(load_checkpoint)[0]
model_dict = paddle.load(load_checkpoint)
self.set_dict(model_dict)
print("load custom checkpoint success")
else:
......@@ -181,7 +181,8 @@ class UserGuidedColorization(nn.Layer):
if not os.path.exists(checkpoint):
os.system('wget https://paddlehub.bj.bcebos.com/dygraph/image_colorization/user_guided.pdparams -O ' +
checkpoint)
model_dict = paddle.load(checkpoint)[0]
model_dict = paddle.load(checkpoint)
self.set_dict(model_dict)
print("load pretrained checkpoint success")
......
hub_module/modules/image/style_transfer/msgnet/module.py
浏览文件 @ 821d6c14
......@@ -5,6 +5,7 @@ import paddle.nn as nn
import numpy as np
import paddle.nn.functional as F
from paddlehub.env import MODULE_HOME
from paddlehub.module.module import moduleinfo
from paddlehub.process.transforms import Compose, Resize, CenterCrop, SetType
from paddlehub.module.cv_module import StyleTransferModule
......@@ -13,7 +14,7 @@ from paddlehub.module.cv_module import StyleTransferModule
class GramMatrix(nn.Layer):
"""Calculate gram matrix"""
def forward(self, y):
(b, ch, h, w) = y.size()
(b, ch, h, w) = y.shape
features = y.reshape((b, ch, w * h))
features_t = features.transpose((0, 2, 1))
gram = features.bmm(features_t) / (ch * h * w)
......@@ -25,8 +26,8 @@ class ConvLayer(nn.Layer):
def __init__(self, in_channels: int, out_channels: int, kernel_size: int, stride: int):
super(ConvLayer, self).__init__()
pad = int(np.floor(kernel_size / 2))
self.reflection_pad = nn.ReflectionPad2d([pad, pad, pad, pad])
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
self.reflection_pad = nn.Pad2D([pad, pad, pad, pad], mode='reflect')
self.conv2d = nn.Conv2D(in_channels, out_channels, kernel_size, stride)
def forward(self, x: paddle.Tensor):
out = self.reflection_pad(x)
......@@ -53,11 +54,11 @@ class UpsampleConvLayer(nn.Layer):
super(UpsampleConvLayer, self).__init__()
self.upsample = upsample
if upsample:
self.upsample_layer = nn.UpSample(scale_factor=upsample)
self.upsample_layer = nn.Upsample(scale_factor=upsample)
self.pad = int(np.floor(kernel_size / 2))
if self.pad != 0:
self.reflection_pad = nn.ReflectionPad2d([self.pad, self.pad, self.pad, self.pad])
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
self.reflection_pad = nn.Pad2D([self.pad, self.pad, self.pad, self.pad], mode='reflect')
self.conv2d = nn.Conv2D(in_channels, out_channels, kernel_size, stride)
def forward(self, x):
if self.upsample:
......@@ -78,7 +79,7 @@ class Bottleneck(nn.Layer):
planes(int): Number of output channels.
stride(int): Number of stride.
downsample(int): Scale factor for downsample layer, default is None.
norm_layer(nn.Layer): Batch norm layer, default is nn.BatchNorm2d.
norm_layer(nn.Layer): Batch norm layer, default is nn.BatchNorm2D.
Return:
img(paddle.Tensor): Bottleneck output.
......@@ -88,18 +89,16 @@ class Bottleneck(nn.Layer):
planes: int,
stride: int = 1,
downsample: int = None,
norm_layer: nn.Layer = nn.BatchNorm2d):
norm_layer: nn.Layer = nn.BatchNorm2D):
super(Bottleneck, self).__init__()
self.expansion = 4
self.downsample = downsample
if self.downsample is not None:
self.residual_layer = nn.Conv2d(inplanes, planes * self.expansion, kernel_size=1, stride=stride)
conv_block = (norm_layer(inplanes), nn.ReLU(), nn.Conv2d(inplanes, planes, kernel_size=1, stride=1),
self.residual_layer = nn.Conv2D(inplanes, planes * self.expansion, kernel_size=1, stride=stride)
conv_block = (norm_layer(inplanes), nn.ReLU(), nn.Conv2D(inplanes, planes, kernel_size=1, stride=1),
norm_layer(planes), nn.ReLU(), ConvLayer(planes, planes, kernel_size=3, stride=stride),
norm_layer(planes), nn.ReLU(), nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
norm_layer(planes), nn.ReLU(), nn.Conv2D(planes, planes * self.expansion, kernel_size=1,
stride=1))
self.conv_block = nn.Sequential(*conv_block)
def forward(self, x: paddle.Tensor):
......@@ -120,12 +119,12 @@ class UpBottleneck(nn.Layer):
inplanes(int): Number of input channels.
planes(int): Number of output channels.
stride(int): Number of stride, default is 2.
norm_layer(nn.Layer): Batch norm layer, default is nn.BatchNorm2d.
norm_layer(nn.Layer): Batch norm layer, default is nn.BatchNorm2D.
Return:
img(paddle.Tensor): UpBottleneck output.
"""
def __init__(self, inplanes: int, planes: int, stride: int = 2, norm_layer: nn.Layer = nn.BatchNorm2d):
def __init__(self, inplanes: int, planes: int, stride: int = 2, norm_layer: nn.Layer = nn.BatchNorm2D):
super(UpBottleneck, self).__init__()
self.expansion = 4
self.residual_layer = UpsampleConvLayer(inplanes,
......@@ -134,20 +133,17 @@ class UpBottleneck(nn.Layer):
stride=1,
upsample=stride)
conv_block = []
conv_block += [norm_layer(inplanes), nn.ReLU(), nn.Conv2d(inplanes, planes, kernel_size=1, stride=1)]
conv_block += [norm_layer(inplanes), nn.ReLU(), nn.Conv2D(inplanes, planes, kernel_size=1, stride=1)]
conv_block += [
norm_layer(planes),
nn.ReLU(),
UpsampleConvLayer(planes, planes, kernel_size=3, stride=1, upsample=stride)
]
conv_block += [
norm_layer(planes),
nn.ReLU(),
nn.Conv2d(planes, planes * self.expansion, kernel_size=1, stride=1)
nn.Conv2D(planes, planes * self.expansion, kernel_size=1, stride=1)
]
self.conv_block = nn.Sequential(*conv_block)
def forward(self, x: paddle.Tensor):
......@@ -195,29 +191,29 @@ class Vgg16(nn.Layer):
""" First four layers from Vgg16."""
def __init__(self):
super(Vgg16, self).__init__()
self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)
self.conv1_1 = nn.Conv2D(3, 64, kernel_size=3, stride=1, padding=1)
self.conv1_2 = nn.Conv2D(64, 64, kernel_size=3, stride=1, padding=1)
self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)
self.conv2_1 = nn.Conv2D(64, 128, kernel_size=3, stride=1, padding=1)
self.conv2_2 = nn.Conv2D(128, 128, kernel_size=3, stride=1, padding=1)
self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1)
self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)
self.conv3_1 = nn.Conv2D(128, 256, kernel_size=3, stride=1, padding=1)
self.conv3_2 = nn.Conv2D(256, 256, kernel_size=3, stride=1, padding=1)
self.conv3_3 = nn.Conv2D(256, 256, kernel_size=3, stride=1, padding=1)
self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1)
self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.conv4_1 = nn.Conv2D(256, 512, kernel_size=3, stride=1, padding=1)
self.conv4_2 = nn.Conv2D(512, 512, kernel_size=3, stride=1, padding=1)
self.conv4_3 = nn.Conv2D(512, 512, kernel_size=3, stride=1, padding=1)
self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)
self.conv5_1 = nn.Conv2D(512, 512, kernel_size=3, stride=1, padding=1)
self.conv5_2 = nn.Conv2D(512, 512, kernel_size=3, stride=1, padding=1)
self.conv5_3 = nn.Conv2D(512, 512, kernel_size=3, stride=1, padding=1)
checkpoint = os.path.join(self.directory, 'vgg16.pdparams')
checkpoint = os.path.join(MODULE_HOME, 'msgnet', 'vgg16.pdparams')
if not os.path.exists(checkpoint):
os.system('wget https://bj.bcebos.com/paddlehub/model/image/image_editing/vgg_paddle.pdparams -O ' +
checkpoint)
model_dict = paddle.load(checkpoint)[0]
model_dict = paddle.load(checkpoint)
self.set_dict(model_dict)
print("load pretrained vgg16 checkpoint success")
......@@ -249,7 +245,7 @@ class Vgg16(nn.Layer):
@moduleinfo(
name="msgnet",
type="CV/image_editing",
author="paddlepaddle",
author="baidu-vis",
author_email="",
summary="Msgnet is a image colorization style transfer model, this module is trained with COCO2014 dataset.",
version="1.0.0",
......@@ -264,7 +260,7 @@ class MSGNet(nn.Layer):
output_nc(int): Number of output channels, default is 3.
ngf(int): Number of input channel for middle layer, default is 128.
n_blocks(int): Block number, default is 6.
norm_layer(nn.Layer): Batch norm layer, default is nn.InstanceNorm2d.
norm_layer(nn.Layer): Batch norm layer, default is nn.InstanceNorm2D.
load_checkpoint(str): Pretrained checkpoint path, default is None.
Return:
......@@ -275,7 +271,7 @@ class MSGNet(nn.Layer):
output_nc=3,
ngf=128,
n_blocks=6,
norm_layer=nn.InstanceNorm2d,
norm_layer=nn.InstanceNorm2D,
load_checkpoint=None):
super(MSGNet, self).__init__()
self.gram = GramMatrix()
......@@ -312,7 +308,7 @@ class MSGNet(nn.Layer):
self.model = nn.Sequential(*model)
if load_checkpoint is not None:
model_dict = paddle.load(load_checkpoint)[0]
model_dict = paddle.load(load_checkpoint)
self.set_dict(model_dict)
print("load custom checkpoint success")
......@@ -321,7 +317,7 @@ class MSGNet(nn.Layer):
if not os.path.exists(checkpoint):
os.system('wget https://bj.bcebos.com/paddlehub/model/image/image_editing/style_paddle.pdparams -O ' +
checkpoint)
model_dict = paddle.load(checkpoint)[0]
model_dict = paddle.load(checkpoint)
model_dict_clone = model_dict.copy()
for key, value in model_dict_clone.items():
if key.endswith(("scale")):
......
paddlehub/module/cv_module.py
浏览文件 @ 821d6c14
......@@ -149,7 +149,7 @@ class ImageColorizeModule(RunModule, ImageServing):
mse = np.mean((visual_ret['real'] * 1.0 - visual_ret['fake_reg'] * 1.0)**2)
psnr_value = 20 * np.log10(255. / np.sqrt(mse))
psnrs.append(psnr_value)
psnr = paddle.to_variable(np.array(psnrs))
psnr = paddle.to_tensor(np.array(psnrs))
return {'loss': loss, 'metrics': {'psnr': psnr}}
def predict(self, images: str, visualization: bool = True, save_path: str = 'result'):
......@@ -239,17 +239,16 @@ class Yolov3Module(RunModule, ImageServing):
for i, out in enumerate(outputs):
anchor_mask = self.anchor_masks[i]
loss = F.yolov3_loss(
x=out,
gt_box=gtbox,
gt_label=gtlabel,
gt_score=gtscore,
anchors=self.anchors,
anchor_mask=anchor_mask,
class_num=self.class_num,
ignore_thresh=self.ignore_thresh,
downsample_ratio=32,
use_label_smooth=False)
loss = F.yolov3_loss(x=out,
gt_box=gtbox,
gt_label=gtlabel,
gt_score=gtscore,
anchors=self.anchors,
anchor_mask=anchor_mask,
class_num=self.class_num,
ignore_thresh=self.ignore_thresh,
downsample_ratio=32,
use_label_smooth=False)
losses.append(paddle.mean(loss))
self.downsample //= 2
......@@ -288,14 +287,13 @@ class Yolov3Module(RunModule, ImageServing):
mask_anchors.append((self.anchors[2 * m]))
mask_anchors.append(self.anchors[2 * m + 1])
box, score = F.yolo_box(
x=out,
img_size=im_shape,
anchors=mask_anchors,
class_num=self.class_num,
conf_thresh=self.valid_thresh,
downsample_ratio=self.downsample,
name="yolo_box" + str(i))
box, score = F.yolo_box(x=out,
img_size=im_shape,
anchors=mask_anchors,
class_num=self.class_num,
conf_thresh=self.valid_thresh,
downsample_ratio=self.downsample,
name="yolo_box" + str(i))
boxes.append(box)
scores.append(paddle.transpose(score, perm=[0, 2, 1]))
......@@ -304,14 +302,13 @@ class Yolov3Module(RunModule, ImageServing):
yolo_boxes = paddle.concat(boxes, axis=1)
yolo_scores = paddle.concat(scores, axis=2)
pred = F.multiclass_nms(
bboxes=yolo_boxes,
scores=yolo_scores,
score_threshold=self.valid_thresh,
nms_top_k=self.nms_topk,
keep_top_k=self.nms_posk,
nms_threshold=self.nms_thresh,
background_label=-1)
pred = F.multiclass_nms(bboxes=yolo_boxes,
scores=yolo_scores,
score_threshold=self.valid_thresh,
nms_top_k=self.nms_topk,
keep_top_k=self.nms_posk,
nms_threshold=self.nms_thresh,
background_label=-1)
bboxes = pred.numpy()
labels = bboxes[:, 0].astype('int32')
......
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或