Im generating scenes in blender, then making a 'realistic' render and a mask render for each scene. To get the mask render Im setting transparent everything that should be and changing colors of objects that I need masks for to shades of gray

Example:

Here are the functions for setting materials and rendering

def set_material_for_objects(objects):
    if len(objects) > 0:
        num_materials = len(objects)
        for i, obj in enumerate(objects):
            grayscale = 0.5 + i * 0.1

            # Create a new material
            material_name = "Material_{:02d}".format(i)
            temp_material = bpy.data.materials.new(name=material_name)
            temp_material.use_nodes = True

            # Get the material output node
            output_node = temp_material.node_tree.nodes.get("Material Output")
            
            # Create an emission node and set its color to the grayscale value
            emission_node = temp_material.node_tree.nodes.new("ShaderNodeEmission")
            emission_node.inputs[0].default_value = (grayscale, grayscale, grayscale, 1.0)
            emission_node.inputs[1].default_value = 1
            
            # Connect the emission node to the material output node
            temp_material.node_tree.links.new(emission_node.outputs[0], output_node.inputs[0])
            
            
            # Assign the material to the object
            obj.data.materials.clear()
            obj.data.materials.append(temp_material)

in short Im using ShaderNodeEmission and setting its color to a shade of gray

def render_scene(objects, visible_objects, box_obj, i):
    
    curr_folder_name = '/IMG' + f'{i:03}'
    
    # RAW RENDER
    bpy.context.scene.render.filepath = FOLDER_PATH + curr_folder_name + "/IMG" + f'{i:03}_render' + ".png"
    bpy.ops.render.render(write_still=True)
    
    # LABEL RENDER
    # set up render for label

    set_transparent([box_obj])  # set the box to be transparent
    set_transparent(objects)     # set objects to be transparent
    set_opaque(visible_objects)  # set visible objects opaque
    set_material_for_objects(visible_objects)  # set visible objects to shades of gray
    
    light = bpy.data.objects['Light']
    light.hide_render = True
    
    # Store the original display device
    original_display_device = bpy.context.scene.display_settings.display_device
    original_aa_settings = bpy.context.scene.view_layers[0].cycles.use_denoising

    # Set up render settings for label
    bpy.context.scene.display_settings.display_device = 'None'
    bpy.context.scene.view_layers[0].cycles.use_denoising = False
    bpy.context.scene.view_layers["ViewLayer"].use_pass_diffuse_color = True

    # render for label
    bpy.context.scene.render.filepath = FOLDER_PATH + curr_folder_name + "/IMG" + f'{i:03}_label' + ".png"
    bpy.ops.render.render(write_still=True)
    
    # return to settings for raw
    set_opaque([box_obj])  # return the box to be opaque
    set_opaque(objects)
    light.hide_render = False
    
    # Set the display device back to the original value ("sRGB" in this case)
    bpy.context.scene.display_settings.display_device = original_display_device
    bpy.context.scene.view_layers[0].cycles.use_denoising = original_aa_settings
    bpy.context.scene.view_layers["ViewLayer"].use_pass_diffuse_color = False
    
    return i + 1

Here I first render and save without any trickery and then render and save only diffuse pass while setting transparent objects that are supposed to be.

full disclosure, I abused chat GPT to write the parts that interact with bpy module as I never worked with it before and didnt have time to properly study it, so it just might be gibberish

So, the question is: how to make these masks have exactly one color? with a naked eye I can alredy see some gradients along the edges. With further inspection the masks are completly out of wack:

mask = Image.open('label2.png').convert('P')
# each mask instance has a different color, from zero to N, where
# N is the number of instances. In order to make visualization easier,
# let's adda color palette to the mask.
mask.putpalette([
    0, 0, 0, # black background
    255, 0, 0, # index 1 is red
    0, 255, 0, # index 2 is yellow
    0, 0, 255, # index 3 is blue
    255, 255, 0, # spare indices
    0,255, 255,
    255,100, 100,
    100, 255, 255,
])

mask

It is also possible that the masks are the way I described I want them to be and Im just cheking the wrong way