Beta

Style transfer is currently in Beta. The API and functionality may change without warning in future releases. More details.

Style Transfer¶

The Task¶

The Neural Style Transfer Task is an optimization method which extract the style from an image and apply it another image while preserving its content. The goal is that the output image looks like the content image, but “painted” in the style of the style reference image.

The StyleTransfer and StyleTransferData classes internally rely on pystiche.

Example¶

Let’s look at transferring the style from The Starry Night onto the images from the COCO 128 data set from the Object Detection Guide. Once we’ve downloaded the data using download_data(), we create the StyleTransferData. Next, we create our StyleTransfer task with the desired style image and fit on the COCO 128 images. We then use the trained StyleTransfer for inference. Finally, we save the model. Here’s the full example:

import os

import flash
import torch
from flash.core.data.utils import download_data
from flash.image.style_transfer import StyleTransfer, StyleTransferData

# 1. Create the DataModule
download_data("https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip", "./data")

datamodule = StyleTransferData.from_folders(train_folder="data/coco128/images/train2017", batch_size=1)

# 2. Build the task
model = StyleTransfer(os.path.join(flash.ASSETS_ROOT, "starry_night.jpg"))

# 3. Create the trainer and train the model
trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count())
trainer.fit(model, datamodule=datamodule)

# 4. Apply style transfer to a few images!
datamodule = StyleTransferData.from_files(
    predict_files=[
        "data/coco128/images/train2017/000000000625.jpg",
        "data/coco128/images/train2017/000000000626.jpg",
        "data/coco128/images/train2017/000000000629.jpg",
    ],
    batch_size=3,
)
predictions = trainer.predict(model, datamodule=datamodule)
print(predictions)

# 5. Save the model!
trainer.save_checkpoint("style_transfer_model.pt")

To learn how to view the available backbones / heads for this task, see Backbones and Heads.

Flash Zero¶

The style transfer task can be used directly from the command line with zero code using Flash Zero. You can run the above example with:

flash style_transfer

To view configuration options and options for running the style transfer task with your own data, use:

flash style_transfer --help