利用GPU Picking实现高性能物体拾取

import * as THREE from 'three'
import { OrbitControls } from 'three/addons/controls/OrbitControls.js'

// --- 0. 关键配置：关闭色彩管理 ---
// 防止 ID 颜色被 Gamma 校正导致数值偏差
THREE.ColorManagement.enabled = false

// --- 1. 初始化场景 ---
const width = window.innerWidth
const height = window.innerHeight
const canvas = document.createElement('canvas')
document.body.appendChild(canvas)

const renderer = new THREE.WebGLRenderer({ canvas: canvas, antialias: true })
renderer.setSize(width, height)
renderer.setPixelRatio(window.devicePixelRatio)

// 创建两个场景：一个用于显示，一个用于拾取
const scene = new THREE.Scene()
scene.background = new THREE.Color(0x111111)

const pickingScene = new THREE.Scene()
pickingScene.background = new THREE.Color(0x000000) // 拾取场景背景必须是纯黑(ID:0)

const camera = new THREE.PerspectiveCamera(75, width / height, 0.1, 1000)
camera.position.set(20, 20, 20)

const controls = new OrbitControls(camera, canvas)
controls.enableDamping = true

// 灯光
scene.add(new THREE.AmbientLight(0xffffff, 0.6))
const sun = new THREE.PointLight(0xffffff, 1000)
sun.position.set(20, 20, 20)
scene.add(sun)

// --- 2. 创建 InstancedMesh ---
const count = 5000
const geometry = new THREE.BoxGeometry(0.5, 0.5, 0.5)

// 视觉 Mesh
const material = new THREE.MeshPhongMaterial()
const mesh = new THREE.InstancedMesh(geometry, material, count)
scene.add(mesh)

// 拾取 Mesh (使用 BasicMaterial，无光照)
const pickingMaterial = new THREE.MeshBasicMaterial()
const pickingMesh = new THREE.InstancedMesh(geometry, pickingMaterial, count)
pickingScene.add(pickingMesh)

// 高亮 Mesh (Wireframe)
const highlightMesh = new THREE.Mesh(
		new THREE.BoxGeometry(0.55, 0.55, 0.55),
		new THREE.MeshBasicMaterial({ color: 0xffff00, wireframe: true })
)
highlightMesh.visible = false
scene.add(highlightMesh)

// --- 3. 填充数据 ---
const matrix = new THREE.Matrix4()
const color = new THREE.Color()
const pickingColor = new THREE.Color()
const matrices = [] // 保存矩阵用于高亮定位

for (let i = 0; i < count; i++) {
		// 随机变换矩阵
		randomizeMatrix(matrix)
		matrices.push(matrix.clone()) // 存起来

		// 同时设置给 视觉Mesh 和 拾取Mesh
		mesh.setMatrixAt(i, matrix)
		pickingMesh.setMatrixAt(i, matrix)

		// 设置视觉颜色
		mesh.setColorAt(i, color.setHSL(Math.random(), 0.7, 0.5))

		// 设置 ID 颜色 (ID = i + 1)
		pickingMesh.setColorAt(i, pickingColor.setHex(i + 1))
}

function randomizeMatrix(matrix) {
		const position = new THREE.Vector3(
				(Math.random() - 0.5) * 30,
				(Math.random() - 0.5) * 30,
				(Math.random() - 0.5) * 30
		)
		const rotation = new THREE.Euler(
				Math.random() * 2 * Math.PI,
				Math.random() * 2 * Math.PI,
				Math.random() * 2 * Math.PI
		)
		const quaternion = new THREE.Quaternion().setFromEuler(rotation)
		const scale = new THREE.Vector3(1, 1, 1)
		matrix.compose(position, quaternion, scale)
}

// --- 4. GPU Picking 核心逻辑 (参考 main.ts) ---

// 创建 1x1 的 RenderTarget，用于读取单像素
const pickingRenderTarget = new THREE.WebGLRenderTarget(1, 1, {
		minFilter: THREE.NearestFilter,
		magFilter: THREE.NearestFilter,
		format: THREE.RGBAFormat
})
const pixelBuffer = new Uint8Array(4)

// 记录鼠标坐标
const mouse = new THREE.Vector2(-1, -1)

function pick() {
		if (mouse.x <= 0 || mouse.y <= 0) return

		// A. 设置相机视口偏移 (仅渲染鼠标下方的 1x1 像素)
		// setViewOffset(fullWidth, fullHeight, x, y, width, height)
		camera.setViewOffset(
				renderer.domElement.width, 
				renderer.domElement.height,
				mouse.x * window.devicePixelRatio, 
				mouse.y * window.devicePixelRatio,
				1, 
				1
		)

		// B. 渲染 PickingScene 到 RenderTarget
		renderer.setRenderTarget(pickingRenderTarget)
		renderer.render(pickingScene, camera)

		// C. 读取像素
		renderer.readRenderTargetPixels(
				pickingRenderTarget,
				0, 0, 1, 1,
				pixelBuffer
		)

		// D. 恢复渲染状态
		renderer.setRenderTarget(null)
		camera.clearViewOffset() // 重要：清除偏移，否则主场景渲染会乱

		// E. 解码 ID
		const id = (pixelBuffer[0] << 16) | (pixelBuffer[1] << 8) | pixelBuffer[2]
		const instanceId = id - 1

		// F. 高亮逻辑
		if (instanceId >= 0) {
				const targetMatrix = matrices[instanceId]
				// 分解矩阵应用到高亮 Mesh
				targetMatrix.decompose(
						highlightMesh.position,
						highlightMesh.quaternion,
						highlightMesh.scale
				)
				highlightMesh.visible = true
				document.body.style.cursor = 'pointer'
		} else {
				highlightMesh.visible = false
				document.body.style.cursor = 'default'
		}
}

// --- 5. 事件监听 ---
window.addEventListener('mousemove', (e) => {
		const rect = canvas.getBoundingClientRect()
		mouse.x = e.clientX - rect.left
		mouse.y = e.clientY - rect.top
})

window.addEventListener('resize', () => {
		const width = window.innerWidth
		const height = window.innerHeight
		camera.aspect = width / height
		camera.updateProjectionMatrix()
		renderer.setSize(width, height)
})

// --- 6. 动画循环 ---
function animate() {
	requestAnimationFrame(animate)
	controls.update()

	// 执行拾取
	pick()

	// 渲染主场景
	renderer.render(scene, camera)
}
animate()