Custom integrator that gets the shape_id through refracting material

[erikbwu]

Hello everyone 👋

I am trying to get the shape_id of a scene for semantic segmentation masks similar to this issue #955 .

However, instead of always returning the shape_id of the first intersection, I want to refract the ray if it hits a specific refracting material (e.g. elm__7 and lets assume it is always this shape string). I tried to use the approach from writing a custom bsdf plugin, but this still gave me very noisy masks, although the integrator is called with spp=1 and also the camera has the filter type box.
I also tried using using the bsdf.sample() to get the new ray, but I always want to refract the ray and not reflect it, and in my case the function returns a zero spectrum

This is the code that I wrote. If you could point me to some documentation or if you have any tips on how to solve this, I would be very thankful.

import mitsuba as mi
import drjit as dr


def get_shape_color(shape, id_to_color):
    """Returns the color of the shape based on its ID or white."""
    return id_to_color.get(shape.id(), mi.Color3f(1, 1, 1))


def process_refract_interaction(pi: mi.PreliminaryIntersection3f, ray, scene):
    """Handles the interaction with the material and computes the refracted ray."""
    si = pi.compute_surface_interaction(ray)
    normal = si.n  # Surface normal
    wi = si.wi  # Incident direction
    cos_theta_i = mi.Frame3f(normal).cos_theta(wi)
    #int_ior = si.shape.bsdf().eval_attribute("int_ior", si), is there a way to get the int_ior attributes directly?

    ext_ior = 1.0   #air
    int_ior = 1.490 #glass
    eta = ext_ior / int_ior  # Refraction index of the lens material
    F, cos_theta_t, eta_it, eta_ti = mi.fresnel(cos_theta_i, eta)

    #refracted_ray = mi.refract(wi, normal, cos_theta_t, eta_ti)  not sure which one to use
    refracted_ray = mi.refract(wi, cos_theta_t, eta_ti)

    # Create a new ray starting from the intersection point
    new_ray = si.spawn_ray(refracted_ray)

    return new_ray

def is_target_shape(shape, target_id):
    """Returns True if the shape ID matches the target ID."""
    return dr.select(shape.id() == "elm__7", dr.cuda.UInt(1), dr.cuda.UInt(0))

class IdIntegrator(mi.SamplingIntegrator):

    def __init__(self, props=mi.Properties()):
        super().__init__(props)

        ids_as_str = props.get("ids_as_str")
        self._id_idx = ids_as_str.split(",")

        self.colormap = [
            mi.Color3f(0.1216, 0.4667, 0.7059),
            mi.Color3f(1.0000, 0.4980, 0.0549),
            mi.Color3f(0.1725, 0.6275, 0.1725),
            mi.Color3f(0.8392, 0.1529, 0.1569),
            mi.Color3f(0.5804, 0.4039, 0.7412),
            mi.Color3f(0.5490, 0.3373, 0.2941),
            mi.Color3f(0.8902, 0.4667, 0.7608),
            mi.Color3f(0.4980, 0.4980, 0.4980),
            mi.Color3f(0.7373, 0.7412, 0.1333),
            mi.Color3f(0.0902, 0.7451, 0.8118),
        ]

        self.id_to_color = {
            id_str: self.colormap[idx % len(self.colormap)] for idx, id_str in enumerate(self._id_idx)
        }
        print(f"ID to color: {self.id_to_color}")

    def sample(
            self,
            scene: mi.Scene,
            sampler: mi.Sampler,
            ray: mi.RayDifferential3f,
            medium: mi.Medium = None,
            active: bool = True,
    ) -> tuple[mi.Color3f, bool, list[float]]:

        def refract_from_material(pi: mi.PreliminaryIntersection3f, ray: mi.RayDifferential3f):
            ray = process_refract_interaction(pi, ray, scene)
            pi = scene.ray_intersect_preliminary(ray)
            return pi

        ray = mi.Ray3f(ray)

        pi: mi.PreliminaryIntersection3f = scene.ray_intersect_preliminary(ray)
        target_id = "elm__7"
        indices = dr.dispatch(pi.shape, is_target_shape, target_id) # target_id is not used, but somehow cannot use dispatch without argument
        pi = dr.switch(indices, [lambda pi, x: pi , refract_from_material], pi, ray)

        shape_color = dr.dispatch(pi.shape, get_shape_color, self.id_to_color)

        # Use the shape color for valid intersections, black otherwise
        mask = dr.select(pi.is_valid(), shape_color, mi.Color3f(0.0))
        return mask, True, []


mi.register_integrator(
    "id_integrator", lambda props: IdIntegrator(props)
)

[njroussel]

Hi @erikbwu

A few of different things:

1. I'm not sure I see why the output is noisy. Could you share a full script, please? You can use the mi.cornell_box() as a dummy scene.

2. You should be able to use the bsdf.sample method to specify what you want. The method takes a BSDFContext (documentation) which allows you to target what types of interaction you want (for a dielectic BSDF that would be ctx.component = 1).

3. The extra argument in the dr.dispatch call should not be necessary. I wonder if it's because you're explicitly using self instead of shape in your is_target_method definition. Have a look at the examples in the documentation to see how these are usually defined.