Gradient masks

If you want to folllow along on google colab, check out this notebook.

Torch dreams can be used to apply masks on different custom losses to generate beautiful blends of channels/layers. It's actually easier done than said.

import cv2
import numpy as np
import matplotlib.pyplot as plt
from torch_dreams.dreamer import dreamer
import torchvision.models as models

First off, we genrate some random noise to work on. Feel free to use any other image in your case.

image_sample = (np.random.rand(512,512*6,3).astype(np.float32) *255).astype(np.uint8)
cv2.imwrite("noise.jpg", image_sample)

Then, we generate the masks to be used using numpy.

def roll(im, shift):
    im = np.roll(im, shift, axis = 1)
    return im
    
grad_mask =  1- np.abs(1- cv2.rotate(np.repeat(np.linspace(2, 0, 512),512*3).reshape(1024,256,3).astype(np.float32) ,cv2.ROTATE_90_CLOCKWISE))
grad_mask = cv2.resize(grad_mask, (512*6, 512))**3
grad_mask_l = roll(grad_mask, -300)
grad_mask_r = roll(grad_mask, 300)
grad_mask_ll = roll(grad_mask, -450*2)
grad_mask_rr = roll(grad_mask, 450*2)
grad_mask_lll =  cv2.resize(cv2.rotate(np.repeat(np.linspace(0, 1, 512)**2,512*3).reshape(1024,256,3).astype(np.float32) ,cv2.ROTATE_90_CLOCKWISE), (512*6, 512))**2.5
grad_mask_rrr =  cv2.resize(cv2.rotate(np.repeat(np.linspace(1, 0, 512)**2,512*3).reshape(1024,256,3).astype(np.float32) ,cv2.ROTATE_90_CLOCKWISE), (512*6, 512))**2.5

grad_masks = [
              grad_mask_lll,
              grad_mask_ll,
              grad_mask_l,
              grad_mask_r,
              grad_mask_rr,
              grad_mask_rrr
]
If things go right, you'll end up with these masks.

Then, we write a helper function to conveniently generate custom loss functions to select certain channels. More on this later.

def make_custom_func(layer_number = 0, channel_number= 0): 
    def custom_func(layer_outputs):
        loss = layer_outputs[layer_number][channel_number].mean()
        return loss
    return custom_func

Now that we're done making the masks, let's fire up out torchdreams.dreamer on the inception_v3 and select a layer to work on. Feel free to choose any other layer(s).

model = models.inception_v3(pretrained=True)
dreamy_boi = dreamer(model)

layers_to_use = [model.Mixed_6c.branch7x7_1.conv]

Now, we generate the custom loss functions required to optimize various chanels, in my case, I only used one layer.

channels = [7,31,115,120,74,7]

custom_funcs = []
for c in channels:
    custom_funcs.append(make_custom_func(channel_number = c))

Then we set up the config

config = {
    "image_path": "noise.jpg",
    "layers": layers_to_use,
    "octave_scale": 1.1,  ## 1.1
    "num_octaves": 20,  ## 14
    "iterations": 100,   ## 100
    "lr": 0.04, ## 0.05
    "max_rotation": 0.7,
    "gradient_smoothing_coeff": None,
    "gradient_smoothing_kernel_size": None,
    "grad_mask": grad_masks,
    "custom_func":  custom_funcs,
    "grayscale": False  ## temporary fix to a bug 
}

All that's left now is to run the optimization (this might take a while depending on the image size)

out = dreamy_boi.deep_dream_with_masks(config)
plt.imshow(out)
plt.show()

Last updated

Was this helpful?