OASIS icon indicating copy to clipboard operation
OASIS copied to clipboard

how to add the local noise?

Open Y-J-HONG opened this issue 2 years ago • 1 comments

when i reproduce the local noise, i cant find the code ,so how should i add the local noise in the segment image

Y-J-HONG avatar May 03 '22 02:05 Y-J-HONG

Hi,

I am assuming you refer to the ability of the model to apply different noise to different part of the image.

For this, you can use the following function which generates 3D noise that is locally different for a chosen class ID per image:

def generate_local_noise(labelmap, z1, z2, class_id = None):
    """Generates a 3D noise tensor consisting of 2 different 1D 
    noise vectors: one for a selected class region and another 
    for every other spatial location. The function returns 
    the concatenation of the labelmap and the 3d noise tensor.

    Args:
        labelmap (torch.tensor): size(batch_size, num_classes, height, width)
        z1 (torch.tensor): size(batch_size, z_dim)
        z2 (torch.tensor): size(batch_size, z_dim)
        class_id (torch.tensor): size(batch_size, 1)
    Returns:
        torch.tensor: size(batch_size, num_classes + z_dim, height, width)
    """ 

    if not class_id:
        # pick a random class per image
        available_classes = labelmap.sum(dim=(2,3))
        class_id = torch.multinomial(available_classes, num_samples=1)
        class_id = class_id.view(class_id.size(0), class_id.size(1),1,1)

    index_map = labelmap.argmax(1, keepdim = True)
    mask = (index_map==class_id).float()
    noise_3d = (1-mask)*z1.view(z1.size(0),z1.size(1),1,1) + mask*z2.view(z2.size(0),z2.size(1),1,1)
    seg_plus_3D_noise = torch.cat((noise_3d, labelmap), dim = 1)

    return seg_plus_3D_noise

You can pass the output of this function the forward function of the generator using the input variable, i.e. img = generator(seg_plus_3D_noise): https://github.com/boschresearch/OASIS/blob/39fe99b7fa492acfa1b66ca64a17fd535f9c6b69/models/generator.py#L30

Since you pass the 3D noise directly in this case, you also have to make sure that the following lines are not executed: https://github.com/boschresearch/OASIS/blob/39fe99b7fa492acfa1b66ca64a17fd535f9c6b69/models/generator.py#L34-L39 You can do this by setting opt.no_3d_noise to True, or adding an if-else statement that suits your purpose.

Hope that helps!

edgarschnfld avatar May 03 '22 12:05 edgarschnfld