2.0rc for msgnet and user_guided_colorization

hub_module/modules/image/colorization/user_guided_colorization/module.py

@@ -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

@@ -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

@@ -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')