Repository: Genesis-Embodied-AI/Genesis
File: genesis/ext/pyrender/scene.py (Python)

Code before:
        Returns
        -------
        scene_pr : :class:`Scene`
            A scene containing the same geometry as the trimesh scene.
        """
        # convert trimesh geometries to pyrender geometries
        geometries = {name: Mesh.from_trimesh(geom) for name, geom in trimesh_scene.geometry.items()}

        # create the pyrender scene object
        scene_pr = Scene(bg_color=bg_color, ambient_light=ambient_light)

        # add every node with geometry to the pyrender scene
        for node in trimesh_scene.graph.nodes_geometry:
            pose, geom_name = trimesh_scene.graph[node]
            scene_pr.add(geometries[geom_name], pose=pose)

        return scene_pr

    def sorted_mesh_nodes(self):
        cam_pos = self.get_pose(self.main_camera_node)
        if len(cam_pos.shape) == 3:
            batched_pos = True
            cam_loc = cam_pos[:, :3, 3]
        else:
            batched_pos = False
            cam_loc = cam_pos[:3, 3]
        solid_nodes = []
        trans_nodes = []
        for node in self.mesh_nodes:
            mesh = node.mesh
            if mesh.is_transparent:
                trans_nodes.append(node)
            else:
                solid_nodes.append(node)

        # TODO BETTER SORTING METHOD
        if batched_pos:
            trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc[0]))
            solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc[0]))
        else:
            trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))
            solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))

        return solid_nodes + trans_nodes


Review comment (performance):
What about

```python
cam_pos = self.get_pose(self.main_camera_node)[..., :3, 3]
batched_pos = len(cam_pos.shape) == 2
```

?

Code after:
        Returns
        -------
        scene_pr : :class:`Scene`
            A scene containing the same geometry as the trimesh scene.
        """
        # convert trimesh geometries to pyrender geometries
        geometries = {name: Mesh.from_trimesh(geom) for name, geom in trimesh_scene.geometry.items()}

        # create the pyrender scene object
        scene_pr = Scene(bg_color=bg_color, ambient_light=ambient_light)

        # add every node with geometry to the pyrender scene
        for node in trimesh_scene.graph.nodes_geometry:
            pose, geom_name = trimesh_scene.graph[node]
            scene_pr.add(geometries[geom_name], pose=pose)

        return scene_pr

    def sorted_mesh_nodes(self):
        cam_pos = self.get_pose(self.main_camera_node)
        cam_loc = cam_pos[..., :3, 3]
        batched_pos = len(cam_pos.shape) == 3
        solid_nodes = []
        trans_nodes = []
        for node in self.mesh_nodes:
            mesh = node.mesh
            if mesh.is_transparent:
                trans_nodes.append(node)
            else:
                solid_nodes.append(node)

        # TODO BETTER SORTING METHOD
        if batched_pos:
            # FIXME normally sorting should be done PER scene when having a batched rasterizer render
            trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc[0]))
            solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc[0]))
        else:
            trans_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))
            solid_nodes.sort(key=lambda n: -np.linalg.norm(self.get_pose(n)[:3, 3] - cam_loc))

        return solid_nodes + trans_nodes