boschresearch/OASIS

Create label Map

MuthElizabeth opened this issue · 1 comments

How to create a label map so that it can be given as input to the pre-trained model along with the original image.

Hi,

The input label map should be a one-hot encoding of shape [BxNxWxH], where B is the batch size, N is the number of classes in the dataset, WxH the height and width of images (we used 256x256 for Ade and Coco, 512x256 for cityscapes).

If you want to use your own input, then something like the following code block should be added to test.py:

# --- read label map ---#
label = Image.open("your_label_map.png")
label = TR.functional.resize(label, (256, 256), Image.NEAREST)
label = label.crop((crop_x, crop_y, crop_x + opt.crop_size, crop_y + opt.crop_size))
label = TR.functional.to_tensor(label)
data_i = {"image":image.unsqueeze(0), "label":label.unsqueeze(0)}

# --- preprocess to one-hot and map to gpu ---#
_, label = models.preprocess_input(opt, data_i)

# --- generate and save an image ---#
generated = model(None, label, "generate", None)
image_saver(label, generated, ["ouput.jpg"])

For further details, please refer to dataloaders to see how we pre-process data in dataloaders, and models/models.py - preprocess_input() to see how we create one-hot encodings.