Merge pull request #54 from makepath/feature/object-picking

Add object picking via OptiX single-ray raycasts

Author: brendancol

rtxpy/engine.py

@@ -631,6 +631,17 @@ def fn(v):
             _add_overlay(v, 'aspect', data.data)
         self._submit(fn)
 
+    # ------------------------------------------------------------------
+    # Picking
+    # ------------------------------------------------------------------
+
+    def pick(self, screen_x, screen_y):
+        """Pick geometry at screen coordinates. Returns hit info dict."""
+        def fn(v):
+            origin, direction = v._screen_to_ray(screen_x, screen_y)
+            return v.rtx.pick(origin, direction)
+        return self._submit(fn)
+
     # ------------------------------------------------------------------
     # Layer management
     # ------------------------------------------------------------------
@@ -1487,6 +1498,32 @@ def _get_right(self):
         right = np.cross(world_up, front)
         return right / (np.linalg.norm(right) + 1e-8)
 
+    def _screen_to_ray(self, screen_x, screen_y):
+        """Convert screen pixel coordinates to a world-space ray.
+
+        Returns (origin, direction) as numpy float32 arrays of shape (3,).
+        """
+        front = self._get_front()
+        world_up = np.array([0, 0, 1], dtype=np.float32)
+        right = np.cross(world_up, front)
+        rn = np.linalg.norm(right)
+        if rn > 1e-8:
+            right /= rn
+        else:
+            right = np.array([1, 0, 0], dtype=np.float32)
+        cam_up = np.cross(front, right)
+
+        fov_scale = np.tan(np.radians(self.fov) / 2.0)
+        aspect = self.render_width / max(1, self.render_height)
+
+        # Window coords → NDC  (-1..1)
+        nx = 2.0 * screen_x / max(1, self.width) - 1.0
+        ny = 1.0 - 2.0 * screen_y / max(1, self.height)
+
+        direction = front + nx * fov_scale * aspect * right + ny * fov_scale * cam_up
+        direction = direction / (np.linalg.norm(direction) + 1e-30)
+        return self.position.copy(), direction.astype(np.float32)
+
     def _get_look_at(self):
         """Get the current look-at point."""
         return self.position + self._get_front() * 1000.0
@@ -5029,6 +5066,18 @@ def _handle_mouse_press(self, button, xpos, ypos):
             self._mouse_last_x = xpos
             self._mouse_last_y = ypos
 
+        elif button == 1:  # right click — object picking
+            origin, direction = self._screen_to_ray(xpos, ypos)
+            result = self.rtx.pick(origin, direction)
+            if result['hit']:
+                gid = result['geometry_id'] or '?'
+                px, py, pz = result['position']
+                print(f"Pick: geometry='{gid}'  pos=({px:.1f}, {py:.1f}, {pz:.1f})  "
+                      f"t={result['t']:.1f}  prim={result['primitive_id']}  "
+                      f"instance={result['instance_id']}")
+            else:
+                print("Pick: no geometry hit")
+
     def _handle_mouse_release(self, button):
         """End drag on button release."""
         self._mouse_dragging = False
@@ -5327,6 +5376,7 @@ def _render_help_text(self):
             ("GEOMETRY", [
                 ("N", "Cycle geometry layer"),
                 ("P", "Prev geometry in group"),
+                ("Right-Click", "Pick geometry"),
             ]),
             ("OBSERVERS", [
                 ("1-8", "Select / create observer"),

rtxpy/rtx.py

@@ -1885,6 +1885,60 @@ def trace(self, rays, hits, numRays: int, primitive_ids=None, instance_ids=None,
         return _trace_rays(self._geom_state, rays, hits, numRays, primitive_ids, instance_ids,
                            ray_flags=ray_flags)
 
+    def pick(self, origin, direction) -> dict:
+        """Fire a single ray and return hit info.
+
+        Parameters
+        ----------
+        origin : array-like
+            Ray origin (x, y, z).
+        direction : array-like
+            Ray direction (dx, dy, dz), will be normalized.
+
+        Returns
+        -------
+        dict
+            Keys: 'hit' (bool), 'geometry_id' (str or None),
+            't' (float), 'normal' (tuple), 'position' (tuple),
+            'primitive_id' (int), 'instance_id' (int).
+        """
+        o = np.asarray(origin, dtype=np.float32)
+        d = np.asarray(direction, dtype=np.float32)
+        d = d / (np.linalg.norm(d) + 1e-30)
+
+        rays = cupy.array([o[0], o[1], o[2], 0.001,
+                           d[0], d[1], d[2], 1e10], dtype=cupy.float32)
+        hits = cupy.zeros(4, dtype=cupy.float32)
+        prim_ids = cupy.full(1, -1, dtype=cupy.int32)
+        inst_ids = cupy.full(1, -1, dtype=cupy.int32)
+
+        self.trace(rays, hits, 1, primitive_ids=prim_ids, instance_ids=inst_ids)
+
+        t = float(hits[0])
+        if t > 0:
+            iid = int(inst_ids[0])
+            geom_list = self.list_geometries()
+            geom_id = geom_list[iid] if 0 <= iid < len(geom_list) else None
+            pos = o + d * t
+            return {
+                'hit': True,
+                'geometry_id': geom_id,
+                't': t,
+                'normal': (float(hits[1]), float(hits[2]), float(hits[3])),
+                'position': (float(pos[0]), float(pos[1]), float(pos[2])),
+                'primitive_id': int(prim_ids[0]),
+                'instance_id': iid,
+            }
+        return {
+            'hit': False,
+            'geometry_id': None,
+            't': -1.0,
+            'normal': (0.0, 0.0, 0.0),
+            'position': (0.0, 0.0, 0.0),
+            'primitive_id': -1,
+            'instance_id': -1,
+        }
+
     # -------------------------------------------------------------------------
     # Multi-GAS API
     # -------------------------------------------------------------------------