<p>If you want some ways to preserve non black color in your style transfer model I suggest checking out the github repo <a href="https://github.com/rrmina/neural-style-pytorch" rel="nofollow noreferrer">here</a>. It has .ipnyb notebooks with entire training pipelines, model weights, a good readme, etc. to reference. According to their readme, they try to implement this <a href="https://arxiv.org/pdf/1606.05897.pdf" rel="nofollow noreferrer">paper</a> on preserving color in neural artistic style transfer which should help you. You can also referecne other repos and run some of them on collabs in this paper of codes repo <a href="https://paperswithcode.com/paper/preserving-color-in-neural-artistic-style" rel="nofollow noreferrer">here</a> though I do suggest looking at the first repo first.</p>
<p>If you want to color transfer outside your styles transfer model and rather have two images that transfer color with the help of some functions in a linrary, then I recommend looking at this <a href="https://www.pyimagesearch.com/2014/06/30/super-fast-color-transfer-images/" rel="nofollow noreferrer">tutorial</a></p>
<p>Sarthak Jain</p>
Only Using Colours From A Specific Part of a Picture For Style Transfer
<p>I've got a neural style transfer model. I'm currently working on trying to use different parts of an image to transfer different pictures. I'm wondering how can I get the model to just use the colours present in an image. Below is an example:</p>
<p><a href="https://i.stack.imgur.com/fMtyq.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/fMtyq.png" alt="enter image description here" /></a></p>
<p><a href="https://i.stack.imgur.com/mBHSO.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/mBHSO.jpg" alt="enter image description here" /></a></p>
<p>The picture above is the style image that I have gotten from using thresholding along with the original image. Now the transferred picture is below:</p>
<p><a href="https://i.stack.imgur.com/LQPgs.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/LQPgs.png" alt="enter image description here" /></a></p>
<p>Obviously it's transferred some of the black parts of the image but I only want the non black colours present to be transferred. Below is my code for my model:</p>
<pre><code>import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image
import torchvision.transforms as transforms
import torchvision.models as models
from torchvision.utils import save_image


class VGG(nn.Module):
    def __init__(self):
        super(VGG, self).__init__()

        self.chosen_features = [&quot;0&quot;, &quot;5&quot;, &quot;10&quot;, &quot;19&quot;, &quot;28&quot;]


        self.model = models.vgg19(pretrained=True).features[:29]

    def forward(self, x):
        # Store relevant features
        features = []

        for layer_num, layer in enumerate(self.model):
            x = layer(x)

            if str(layer_num) in self.chosen_features:
                features.append(x)

        return features


device = torch.device(&quot;cuda&quot; if torch.cuda.is_available() else &quot;cpu&quot;)
    
def load_image(image_name):
    image = Image.open(image_name)
    image = loader(image).unsqueeze(0)
    return image.to(device)

imsize = 384

loader = transforms.Compose(
    [
        transforms.Resize((imsize, imsize)),
        transforms.ToTensor(),
        # transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

original_img = load_image(&quot;Content Image.jpg&quot;)
style_img = load_image(&quot;Adaptive Image 2.jpg&quot;)

# initialized generated as white noise or clone of original image.
# Clone seemed to work better for me.

generated = original_img.clone().requires_grad_(True)
# generated = load_image(&quot;20epoctom.png&quot;)
model = VGG().to(device).eval()


# Hyperparameters
total_steps = 10000
learning_rate = 0.001
alpha = 1
beta = 0.01
optimizer = optim.Adam([generated], lr=learning_rate)

for step in range(total_steps):
    # Obtain the convolution features in specifically chosen layers
    generated_features = model(generated)
    original_img_features = model(original_img)
    style_features = model(style_img)

    # Loss is 0 initially
    style_loss = original_loss = 0

    # iterate through all the features for the chosen layers
    for gen_feature, orig_feature, style_feature in zip(
        generated_features, original_img_features, style_features
    ):

        # batch_size will just be 1
        batch_size, channel, height, width = gen_feature.shape
        original_loss += torch.mean((gen_feature - orig_feature) ** 2)
        # Compute Gram Matrix of generated
        G = gen_feature.view(channel, height * width).mm(
            gen_feature.view(channel, height * width).t()
        )
        # Compute Gram Matrix of Style
        A = style_feature.view(channel, height * width).mm(
            style_feature.view(channel, height * width).t()
        )
        style_loss += torch.mean((G - A) ** 2)

    total_loss = alpha * original_loss + beta * style_loss
    optimizer.zero_grad()
    total_loss.backward()
    optimizer.step()
    
    if step % 500 == 0:
        print(total_loss)
        save_image(generated, f&quot;Generated Pictures/{step//500} Iterations Generated Picture.png&quot;)
</code></pre>
<p>Any idea of where to potentially go as well would be appreciated!</p>