Gradient Masks

Gradient masks

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
]

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[0][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). If things go right, you'll end up with a photo similar to the banner in this page.

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