[WIP] add style transfer task with pystiche

flash/utils/imports.py

@@ -8,3 +8,4 @@
 _PYTORCHVIDEO_AVAILABLE = _module_available("pytorchvideo")
 _MATPLOTLIB_AVAILABLE = _module_available("matplotlib")
 _TRANSFORMERS_AVAILABLE = _module_available("transformers")
+_PYSTICHE_AVAILABLE = _module_available("pystiche")

flash/vision/__init__.py

@@ -2,3 +2,4 @@
 from flash.vision.classification import ImageClassificationData, ImageClassificationPreprocess, ImageClassifier
 from flash.vision.detection import ObjectDetectionData, ObjectDetector
 from flash.vision.embedding import ImageEmbedder
+from .style_transfer import *

flash/vision/style_transfer/__init__.py

@@ -0,0 +1,2 @@
+from .data import *
+from .model import *

flash/vision/style_transfer/data.py

@@ -0,0 +1 @@
+# TODO

flash/vision/style_transfer/model.py

@@ -0,0 +1,90 @@
+from torch import nn
+from torch.nn.functional import interpolate
+
+__all__ = ["Transformer"]
+
+
+class Interpolate(nn.Module):
+    def __init__(self, scale_factor=1.0, mode="nearest"):
+        super().__init__()
+        self.scale_factor = scale_factor
+        self.mode = mode
+
+    def forward(self, input):
+        return interpolate(input, scale_factor=self.scale_factor, mode=self.mode)
+
+    def extra_repr(self):
+        extras = []
+        if self.scale_factor:
+            extras.append(f"scale_factor={self.scale_factor}")
+        if self.mode != "nearest":
+            extras.append(f"mode={self.mode}")
+        return ", ".join(extras)
+
+
+class Conv(nn.Module):
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size,
+        stride=1,
+        upsample=False,
+        norm=True,
+        activation=True,
+    ):
+        super().__init__()
+        self.upsample = Interpolate(scale_factor=stride) if upsample else None
+        self.pad = nn.ReflectionPad2d(kernel_size // 2)
+        self.conv = nn.Conv2d(
+            in_channels, out_channels, kernel_size, stride=1 if upsample else stride
+        )
+        self.norm = nn.InstanceNorm2d(out_channels, affine=True) if norm else None
+        self.activation = nn.ReLU() if activation else None
+
+    def forward(self, input):
+        if self.upsample:
+            input = self.upsample(input)
+
+        output = self.conv(self.pad(input))
+
+        if self.norm:
+            output = self.norm(output)
+        if self.activation:
+            output = self.activation(output)
+
+        return output
+
+
+class Residual(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.conv1 = Conv(channels, channels, kernel_size=3)
+        self.conv2 = Conv(channels, channels, kernel_size=3, activation=False)
+
+    def forward(self, input):
+        output = self.conv2(self.conv1(input))
+        return output + input
+
+
+class Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = nn.Sequential(
+            Conv(3, 32, kernel_size=9),
+            Conv(32, 64, kernel_size=3, stride=2),
+            Conv(64, 128, kernel_size=3, stride=2),
+            Residual(128),
+            Residual(128),
+            Residual(128),
+            Residual(128),
+            Residual(128),
+        )
+        self.decoder = nn.Sequential(
+            Conv(128, 64, kernel_size=3, stride=2, upsample=True),
+            Conv(64, 32, kernel_size=3, stride=2, upsample=True),
+            Conv(32, 3, kernel_size=9, norm=False, activation=False),
+        )
+
+    def forward(self, input):
+        return self.decoder(self.encoder(input))

flash_examples/predict/style_transfer.py

@@ -0,0 +1,47 @@
+import sys
+
+import torch
+
+from flash.utils.imports import _PYSTICHE_AVAILABLE
+
+if _PYSTICHE_AVAILABLE:
+    from pystiche import enc, loss, ops
+else:
+    print("Please, run `pip install pystiche`")
+    sys.exit(0)
+
+multi_layer_encoder = enc.vgg16_multi_layer_encoder()
+
+content_layer = "relu2_2"
+content_encoder = multi_layer_encoder.extract_encoder(content_layer)
+content_weight = 1e5
+content_loss = ops.FeatureReconstructionOperator(
+    content_encoder, score_weight=content_weight
+)
+
+
+class GramOperator(ops.GramOperator):
+    def enc_to_repr(self, enc: torch.Tensor) -> torch.Tensor:
+        repr = super().enc_to_repr(enc)
+        num_channels = repr.size()[1]
+        return repr / num_channels
+
+
+style_layers = ("relu1_2", "relu2_2", "relu3_3", "relu4_3")
+style_weight = 1e10
+style_loss = ops.MultiLayerEncodingOperator(
+    multi_layer_encoder,
+    style_layers,
+    lambda encoder, layer_weight: GramOperator(encoder, score_weight=layer_weight),
+    layer_weights="sum",
+    score_weight=style_weight,
+)
+
+# TODO: this needs to be moved to the device to be trained on
+# TODO: we need to register a style image here
+perceptual_loss = loss.PerceptualLoss(content_loss, style_loss)
+
+
+def loss_fn(image):
+    perceptual_loss.set_content_image(image)
+    return float(perceptual_loss(image))

requirements.txt

@@ -19,3 +19,4 @@ pycocotools>=2.0.2 ; python_version >= "3.7"
 kornia==0.5.0
 pytorchvideo
 matplotlib  # used by the visualisation callback
+pystiche>=0.7.1