From 52bad3d0d12c178bf3524d65f3c5532c6412dcff Mon Sep 17 00:00:00 2001
From: Terry Jia <terryjia88@gmail.com>
Date: Tue, 18 Mar 2025 10:51:53 -0400
Subject: [PATCH] [3d] support output normal and lineart at once (#3122)

---
 src/components/load3d/Load3D.vue              |  20 +++-
 src/components/load3d/Load3DAnimation.vue     |  23 +++-
 .../load3d/Load3DAnimationScene.vue           |   4 +
 src/extensions/core/load3d.ts                 |  31 ++++--
 src/extensions/core/load3d/Load3d.ts          |  18 +++
 src/extensions/core/load3d/SceneManager.ts    | 103 +++++++++++++++++-
 src/extensions/core/load3d/interfaces.ts      |   2 +
 src/services/load3dService.ts                 |   6 +
 8 files changed, 194 insertions(+), 13 deletions(-)

diff --git a/src/components/load3d/Load3D.vue b/src/components/load3d/Load3D.vue
index a41743c12..3864822e8 100644
--- a/src/components/load3d/Load3D.vue
+++ b/src/components/load3d/Load3D.vue
@@ -1,6 +1,11 @@
 <template>
-  <div class="relative w-full h-full">
+  <div
+    class="relative w-full h-full"
+    @mouseenter="handleMouseEnter"
+    @mouseleave="handleMouseLeave"
+  >
     <Load3DScene
+      ref="load3DSceneRef"
       :node="node"
       :type="type"
       :backgroundColor="backgroundColor"
@@ -90,11 +95,24 @@ const backgroundImage = ref('')
 const upDirection = ref<UpDirection>('original')
 const materialMode = ref<MaterialMode>('original')
 const edgeThreshold = ref(85)
+const load3DSceneRef = ref(null)
 
 const showPreviewButton = computed(() => {
   return !type.includes('Preview')
 })
 
+const handleMouseEnter = () => {
+  if (load3DSceneRef.value?.load3d) {
+    load3DSceneRef.value.load3d.updateStatusMouseOnScene(true)
+  }
+}
+
+const handleMouseLeave = () => {
+  if (load3DSceneRef.value?.load3d) {
+    load3DSceneRef.value.load3d.updateStatusMouseOnScene(false)
+  }
+}
+
 const switchCamera = () => {
   cameraType.value =
     cameraType.value === 'perspective' ? 'orthographic' : 'perspective'
diff --git a/src/components/load3d/Load3DAnimation.vue b/src/components/load3d/Load3DAnimation.vue
index d181dfbf3..3eb909747 100644
--- a/src/components/load3d/Load3DAnimation.vue
+++ b/src/components/load3d/Load3DAnimation.vue
@@ -1,6 +1,11 @@
 <template>
-  <div class="relative w-full h-full">
+  <div
+    class="relative w-full h-full"
+    @mouseenter="handleMouseEnter"
+    @mouseleave="handleMouseLeave"
+  >
     <Load3DAnimationScene
+      ref="load3DAnimationSceneRef"
       :node="node"
       :type="type"
       :backgroundColor="backgroundColor"
@@ -110,6 +115,22 @@ const showPreviewButton = computed(() => {
   return !type.includes('Preview')
 })
 
+const load3DAnimationSceneRef = ref(null)
+
+const handleMouseEnter = () => {
+  const sceneRef = load3DAnimationSceneRef.value?.load3DSceneRef
+  if (sceneRef?.load3d) {
+    sceneRef.load3d.updateStatusMouseOnScene(true)
+  }
+}
+
+const handleMouseLeave = () => {
+  const sceneRef = load3DAnimationSceneRef.value?.load3DSceneRef
+  if (sceneRef?.load3d) {
+    sceneRef.load3d.updateStatusMouseOnScene(false)
+  }
+}
+
 const switchCamera = () => {
   cameraType.value =
     cameraType.value === 'perspective' ? 'orthographic' : 'perspective'
diff --git a/src/components/load3d/Load3DAnimationScene.vue b/src/components/load3d/Load3DAnimationScene.vue
index cf1c711d4..c6528f7d0 100644
--- a/src/components/load3d/Load3DAnimationScene.vue
+++ b/src/components/load3d/Load3DAnimationScene.vue
@@ -183,4 +183,8 @@ const animationListeners = {
     emit('animationListChange', newValue)
   }
 }
+
+defineExpose({
+  load3DSceneRef
+})
 </script>
diff --git a/src/extensions/core/load3d.ts b/src/extensions/core/load3d.ts
index 2382ceb8b..f79bd680d 100644
--- a/src/extensions/core/load3d.ts
+++ b/src/extensions/core/load3d.ts
@@ -124,21 +124,30 @@ app.registerExtension({
     sceneWidget.serializeValue = async () => {
       node.properties['Camera Info'] = load3d.getCameraState()
 
-      const { scene: imageData, mask: maskData } = await load3d.captureScene(
+      const {
+        scene: imageData,
+        mask: maskData,
+        normal: normalData,
+        lineart: lineartData
+      } = await load3d.captureScene(
         width.value as number,
         height.value as number
       )
 
-      const [data, dataMask] = await Promise.all([
+      const [data, dataMask, dataNormal, dataLineart] = await Promise.all([
         Load3dUtils.uploadTempImage(imageData, 'scene'),
-        Load3dUtils.uploadTempImage(maskData, 'scene_mask')
+        Load3dUtils.uploadTempImage(maskData, 'scene_mask'),
+        Load3dUtils.uploadTempImage(normalData, 'scene_normal'),
+        Load3dUtils.uploadTempImage(lineartData, 'scene_lineart')
       ])
 
       load3d.handleResize()
 
       return {
         image: `threed/${data.name} [temp]`,
-        mask: `threed/${dataMask.name} [temp]`
+        mask: `threed/${dataMask.name} [temp]`,
+        normal: `threed/${dataNormal.name} [temp]`,
+        lineart: `threed/${dataLineart.name} [temp]`
       }
     }
   }
@@ -252,21 +261,27 @@ app.registerExtension({
 
       load3d.toggleAnimation(false)
 
-      const { scene: imageData, mask: maskData } = await load3d.captureScene(
+      const {
+        scene: imageData,
+        mask: maskData,
+        normal: normalData
+      } = await load3d.captureScene(
         width.value as number,
         height.value as number
       )
 
-      const [data, dataMask] = await Promise.all([
+      const [data, dataMask, dataNormal] = await Promise.all([
         Load3dUtils.uploadTempImage(imageData, 'scene'),
-        Load3dUtils.uploadTempImage(maskData, 'scene_mask')
+        Load3dUtils.uploadTempImage(maskData, 'scene_mask'),
+        Load3dUtils.uploadTempImage(normalData, 'scene_normal')
       ])
 
       load3d.handleResize()
 
       return {
         image: `threed/${data.name} [temp]`,
-        mask: `threed/${dataMask.name} [temp]`
+        mask: `threed/${dataMask.name} [temp]`,
+        normal: `threed/${dataNormal.name} [temp]`
       }
     }
   }
diff --git a/src/extensions/core/load3d/Load3d.ts b/src/extensions/core/load3d/Load3d.ts
index 8e4a60984..14f597fbf 100644
--- a/src/extensions/core/load3d/Load3d.ts
+++ b/src/extensions/core/load3d/Load3d.ts
@@ -36,6 +36,9 @@ class Load3d {
   protected loaderManager: LoaderManager
   protected modelManager: ModelManager
 
+  STATUS_MOUSE_ON_NODE: boolean
+  STATUS_MOUSE_ON_SCENE: boolean
+
   constructor(
     container: Element | HTMLElement,
     options: Load3DOptions = {
@@ -124,6 +127,9 @@ class Load3d {
       this.previewManager.init()
     }
 
+    this.STATUS_MOUSE_ON_NODE = false
+    this.STATUS_MOUSE_ON_SCENE = false
+
     this.handleResize()
     this.startAnimation()
   }
@@ -167,6 +173,18 @@ class Load3d {
     animate()
   }
 
+  updateStatusMouseOnNode(onNode: boolean): void {
+    this.STATUS_MOUSE_ON_NODE = onNode
+  }
+
+  updateStatusMouseOnScene(onScene: boolean): void {
+    this.STATUS_MOUSE_ON_SCENE = onScene
+  }
+
+  isActive(): boolean {
+    return this.STATUS_MOUSE_ON_NODE || this.STATUS_MOUSE_ON_SCENE
+  }
+
   setBackgroundColor(color: string): void {
     this.sceneManager.setBackgroundColor(color)
     this.renderer.render(
diff --git a/src/extensions/core/load3d/SceneManager.ts b/src/extensions/core/load3d/SceneManager.ts
index e2a671594..824894802 100644
--- a/src/extensions/core/load3d/SceneManager.ts
+++ b/src/extensions/core/load3d/SceneManager.ts
@@ -196,7 +196,7 @@ export class SceneManager implements SceneManagerInterface {
   captureScene(
     width: number,
     height: number
-  ): Promise<{ scene: string; mask: string }> {
+  ): Promise<{ scene: string; mask: string; normal: string; lineart: string }> {
     return new Promise(async (resolve, reject) => {
       try {
         const originalWidth = this.renderer.domElement.width
@@ -207,6 +207,7 @@ export class SceneManager implements SceneManagerInterface {
         const originalClearAlpha = this.renderer.getClearAlpha()
         const originalToneMapping = this.renderer.toneMapping
         const originalExposure = this.renderer.toneMappingExposure
+        const originalOutputColorSpace = this.renderer.outputColorSpace
 
         this.renderer.setSize(width, height)
 
@@ -240,6 +241,11 @@ export class SceneManager implements SceneManagerInterface {
           )
         }
 
+        const originalMaterials = new Map<
+          THREE.Mesh,
+          THREE.Material | THREE.Material[]
+        >()
+
         this.renderer.clear()
 
         if (this.backgroundMesh && this.backgroundTexture) {
@@ -262,12 +268,103 @@ export class SceneManager implements SceneManagerInterface {
         this.renderer.render(this.scene, this.getActiveCamera())
         const maskData = this.renderer.domElement.toDataURL('image/png')
 
+        this.scene.traverse((child) => {
+          if (child instanceof THREE.Mesh) {
+            originalMaterials.set(child, child.material)
+
+            child.material = new THREE.MeshNormalMaterial({
+              flatShading: false,
+              side: THREE.DoubleSide,
+              normalScale: new THREE.Vector2(1, 1)
+            })
+          }
+        })
+
+        const gridVisible = this.gridHelper.visible
+        this.gridHelper.visible = false
+
+        this.renderer.setClearColor(0x000000, 1)
+        this.renderer.clear()
+        this.renderer.render(this.scene, this.getActiveCamera())
+        const normalData = this.renderer.domElement.toDataURL('image/png')
+
+        this.scene.traverse((child) => {
+          if (child instanceof THREE.Mesh) {
+            const originalMaterial = originalMaterials.get(child)
+            if (originalMaterial) {
+              child.material = originalMaterial
+            }
+          }
+        })
+
+        let lineartModel: THREE.Group | null = null
+
+        const originalSceneVisible: Map<THREE.Object3D, boolean> = new Map()
+
+        this.scene.traverse((child) => {
+          if (child instanceof THREE.Group && child.name === 'lineartModel') {
+            lineartModel = child as THREE.Group
+          }
+
+          if (
+            child instanceof THREE.Mesh &&
+            !(child.parent?.name === 'lineartModel')
+          ) {
+            originalSceneVisible.set(child, child.visible)
+
+            child.visible = false
+          }
+        })
+
+        this.renderer.setClearColor(0xffffff, 1)
+        this.renderer.clear()
+
+        if (lineartModel !== null) {
+          lineartModel = lineartModel as THREE.Group
+
+          const originalLineartVisibleMap: Map<THREE.Object3D, boolean> =
+            new Map()
+
+          lineartModel.traverse((child: THREE.Object3D) => {
+            if (child instanceof THREE.Mesh) {
+              originalLineartVisibleMap.set(child, child.visible)
+
+              child.visible = true
+            }
+          })
+
+          const originalLineartVisible = lineartModel.visible
+          lineartModel.visible = true
+
+          this.renderer.render(this.scene, this.getActiveCamera())
+
+          lineartModel.visible = originalLineartVisible
+
+          originalLineartVisibleMap.forEach((visible, object) => {
+            object.visible = visible
+          })
+        }
+
+        const lineartData = this.renderer.domElement.toDataURL('image/png')
+
+        originalSceneVisible.forEach((visible, object) => {
+          object.visible = visible
+        })
+
+        this.gridHelper.visible = gridVisible
+
         this.renderer.setClearColor(originalClearColor, originalClearAlpha)
         this.renderer.setSize(originalWidth, originalHeight)
+        this.renderer.outputColorSpace = originalOutputColorSpace
 
-        this.handleResize(width, height)
+        this.handleResize(originalWidth, originalHeight)
 
-        resolve({ scene: sceneData, mask: maskData })
+        resolve({
+          scene: sceneData,
+          mask: maskData,
+          normal: normalData,
+          lineart: lineartData
+        })
       } catch (error) {
         reject(error)
       }
diff --git a/src/extensions/core/load3d/interfaces.ts b/src/extensions/core/load3d/interfaces.ts
index cde866638..0c8fd1edc 100644
--- a/src/extensions/core/load3d/interfaces.ts
+++ b/src/extensions/core/load3d/interfaces.ts
@@ -40,6 +40,8 @@ export interface Load3DOptions {
 export interface CaptureResult {
   scene: string
   mask: string
+  normal: string
+  lineart: string
 }
 
 export interface BaseManager {
diff --git a/src/services/load3dService.ts b/src/services/load3dService.ts
index c891e2094..23ca1978d 100644
--- a/src/services/load3dService.ts
+++ b/src/services/load3dService.ts
@@ -41,6 +41,12 @@ export class Load3dService {
 
     rawNode.onMouseEnter = function () {
       instance.refreshViewport()
+
+      instance.updateStatusMouseOnNode(true)
+    }
+
+    rawNode.onMouseLeave = function () {
+      instance.updateStatusMouseOnNode(false)
     }
 
     rawNode.onResize = function () {