OASIS
OASIS copied to clipboard
how to add the local noise?
when i reproduce the local noise, i cant find the code ,so how should i add the local noise in the segment image
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!