Skip to content

Commit c51821d

Browse files
jtydhr88jtydhr88
committed
stl and different material display support
1 parent 793e697 commit c51821d

File tree

1 file changed

+160
-17
lines changed

1 file changed

+160
-17
lines changed

‎src/extensions/core/load3d.ts‎

Lines changed: 160 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
// @ts-strict-ignore
21
import { app } from '@/scripts/app'
32
import { api } from '@/scripts/api'
43
import { useToastStore } from '@/stores/toastStore'
@@ -8,6 +7,7 @@ import { GLTFLoader } from 'three/examples/jsm/loaders/GLTFLoader'
87
import { OBJLoader } from 'three/examples/jsm/loaders/OBJLoader'
98
import { MTLLoader } from 'three/examples/jsm/loaders/MTLLoader'
109
import { FBXLoader } from 'three/examples/jsm/loaders/FBXLoader'
10+
import { STLLoader } from 'three/examples/jsm/loaders/STLLoader'
1111
import { IWidget } from '@comfyorg/litegraph'
1212
import { nextTick } from 'vue'
1313

@@ -21,7 +21,7 @@ async function uploadFile(
2121
try {
2222
const body = new FormData()
2323
body.append('image', file)
24-
body.append('subfolder', 'mesh')
24+
body.append('subfolder', '3d')
2525

2626
const resp = await api.fetchApi('/upload/image', {
2727
method: 'POST',
@@ -53,7 +53,7 @@ async function uploadFile(
5353
if (mtlFile) {
5454
const mtlFormData = new FormData()
5555
mtlFormData.append('image', mtlFile)
56-
mtlFormData.append('subfolder', 'mesh')
56+
mtlFormData.append('subfolder', '3d')
5757

5858
await api.fetchApi('/upload/image', {
5959
method: 'POST',
@@ -86,6 +86,7 @@ class Load3d {
8686
objLoader: OBJLoader
8787
mtlLoader: MTLLoader
8888
fbxLoader: FBXLoader
89+
stlLoader: STLLoader
8990
currentModel: THREE.Object3D | null = null
9091
currentAnimation: THREE.AnimationMixer | null = null
9192
animationActions: THREE.AnimationAction[] = []
@@ -95,6 +96,13 @@ class Load3d {
9596
gridHelper: THREE.GridHelper
9697
lights: THREE.Light[] = []
9798
clock: THREE.Clock
99+
normalMaterial: THREE.MeshNormalMaterial
100+
standardMaterial: THREE.MeshStandardMaterial
101+
wireframeMaterial: THREE.MeshBasicMaterial
102+
originalMaterials: WeakMap<THREE.Mesh, THREE.Material | THREE.Material[]> =
103+
new WeakMap()
104+
105+
materialMode: 'original' | 'normal' | 'wireframe' = 'original'
98106

99107
constructor(container: Element | HTMLElement) {
100108
this.scene = new THREE.Scene()
@@ -136,6 +144,7 @@ class Load3d {
136144
this.objLoader = new OBJLoader()
137145
this.mtlLoader = new MTLLoader()
138146
this.fbxLoader = new FBXLoader()
147+
this.stlLoader = new STLLoader()
139148
this.clock = new THREE.Clock()
140149

141150
this.setupLights()
@@ -144,13 +153,80 @@ class Load3d {
144153
this.gridHelper.position.set(0, 0, 0)
145154
this.scene.add(this.gridHelper)
146155

156+
this.normalMaterial = new THREE.MeshNormalMaterial({
157+
flatShading: false,
158+
side: THREE.DoubleSide,
159+
normalScale: new THREE.Vector2(1, 1),
160+
transparent: false,
161+
opacity: 1.0
162+
})
163+
164+
this.wireframeMaterial = new THREE.MeshBasicMaterial({
165+
color: 0xffffff,
166+
wireframe: true,
167+
transparent: false,
168+
opacity: 1.0
169+
})
170+
171+
this.standardMaterial = this.createSTLMaterial()
172+
147173
this.animate()
148174

149175
this.handleResize()
150176

151177
this.startAnimation()
152178
}
153179

180+
setMaterialMode(mode: 'original' | 'normal' | 'wireframe') {
181+
this.materialMode = mode
182+
183+
if (this.currentModel) {
184+
this.currentModel.traverse((child) => {
185+
if (child instanceof THREE.Mesh) {
186+
switch (mode) {
187+
case 'normal':
188+
if (!this.originalMaterials.has(child)) {
189+
this.originalMaterials.set(child, child.material)
190+
}
191+
child.material = new THREE.MeshNormalMaterial({
192+
flatShading: false,
193+
side: THREE.DoubleSide,
194+
normalScale: new THREE.Vector2(1, 1),
195+
transparent: false,
196+
opacity: 1.0
197+
})
198+
child.geometry.computeVertexNormals()
199+
break
200+
201+
case 'wireframe':
202+
if (!this.originalMaterials.has(child)) {
203+
this.originalMaterials.set(child, child.material)
204+
}
205+
child.material = new THREE.MeshBasicMaterial({
206+
color: 0xffffff,
207+
wireframe: true,
208+
transparent: false,
209+
opacity: 1.0
210+
})
211+
break
212+
213+
case 'original':
214+
const originalMaterial = this.originalMaterials.get(child)
215+
if (originalMaterial) {
216+
child.material = originalMaterial
217+
} else {
218+
child.material = this.standardMaterial
219+
}
220+
break
221+
}
222+
}
223+
})
224+
225+
this.renderer.outputColorSpace = THREE.SRGBColorSpace
226+
this.renderer.render(this.scene, this.activeCamera)
227+
}
228+
}
229+
154230
setupLights() {
155231
const ambientLight = new THREE.AmbientLight(0xffffff, 0.5)
156232
this.scene.add(ambientLight)
@@ -321,6 +397,9 @@ class Load3d {
321397
this.controls.update()
322398

323399
this.renderer.render(this.scene, this.activeCamera)
400+
401+
this.materialMode = 'original'
402+
this.originalMaterials = new WeakMap()
324403
}
325404

326405
toggleAnimation(play?: boolean) {
@@ -358,17 +437,34 @@ class Load3d {
358437

359438
if (!fileExtension) {
360439
useToastStore().addAlert('Could not determine file type')
361-
362440
return
363441
}
364442

365443
let model: THREE.Object3D | null = null
366444

367445
switch (fileExtension) {
446+
case 'stl':
447+
const geometry = await this.stlLoader.loadAsync(url)
448+
geometry.computeVertexNormals()
449+
450+
const mesh = new THREE.Mesh(geometry, this.standardMaterial)
451+
452+
const group = new THREE.Group()
453+
group.add(mesh)
454+
455+
model = group
456+
break
457+
368458
case 'fbx':
369459
const fbxModel = await this.fbxLoader.loadAsync(url)
370460
model = fbxModel
371461

462+
fbxModel.traverse((child) => {
463+
if (child instanceof THREE.Mesh) {
464+
this.originalMaterials.set(child, child.material)
465+
}
466+
})
467+
372468
if (fbxModel.animations.length > 0) {
373469
this.currentAnimation = new THREE.AnimationMixer(fbxModel)
374470
this.animationActions = fbxModel.animations.map((clip) => {
@@ -382,24 +478,39 @@ class Load3d {
382478
break
383479

384480
case 'obj':
385-
const mtlUrl = url.replace(/\.obj([^.]*$)/, '.mtl$1')
386-
try {
387-
const materials = await this.mtlLoader.loadAsync(mtlUrl)
388-
materials.preload()
389-
this.objLoader.setMaterials(materials)
390-
} catch (e) {
391-
console.log(
392-
'No MTL file found or error loading it, continuing without materials'
393-
)
481+
if (this.materialMode === 'original') {
482+
const mtlUrl = url.replace(/\.obj([^.]*$)/, '.mtl$1')
483+
try {
484+
const materials = await this.mtlLoader.loadAsync(mtlUrl)
485+
materials.preload()
486+
this.objLoader.setMaterials(materials)
487+
} catch (e) {
488+
console.log(
489+
'No MTL file found or error loading it, continuing without materials'
490+
)
491+
}
394492
}
395493

396494
model = await this.objLoader.loadAsync(url)
495+
496+
model.traverse((child) => {
497+
if (child instanceof THREE.Mesh) {
498+
this.originalMaterials.set(child, child.material)
499+
}
500+
})
397501
break
398502

399503
case 'gltf':
400504
case 'glb':
401505
const gltf = await this.gltfLoader.loadAsync(url)
402506
model = gltf.scene
507+
508+
gltf.scene.traverse((child) => {
509+
if (child instanceof THREE.Mesh) {
510+
child.geometry.computeVertexNormals()
511+
this.originalMaterials.set(child, child.material)
512+
}
513+
})
403514
break
404515

405516
default:
@@ -427,6 +538,10 @@ class Load3d {
427538

428539
this.scene.add(model)
429540

541+
if (this.materialMode !== 'original') {
542+
this.setMaterialMode(this.materialMode)
543+
}
544+
430545
const distance = Math.max(size.x, size.z) * 2
431546
const height = size.y * 2
432547

@@ -451,6 +566,10 @@ class Load3d {
451566
this.controls.target.set(0, size.y / 2, 0)
452567
this.controls.update()
453568

569+
this.renderer.outputColorSpace = THREE.SRGBColorSpace
570+
this.renderer.toneMapping = THREE.ACESFilmicToneMapping
571+
this.renderer.toneMappingExposure = 1
572+
454573
this.handleResize()
455574
}
456575
} catch (error) {
@@ -532,6 +651,16 @@ class Load3d {
532651
})
533652
}
534653

654+
createSTLMaterial() {
655+
return new THREE.MeshStandardMaterial({
656+
color: 0x808080,
657+
metalness: 0.1,
658+
roughness: 0.8,
659+
flatShading: false,
660+
side: THREE.DoubleSide
661+
})
662+
}
663+
535664
setViewPosition(position: 'front' | 'top' | 'right' | 'isometric') {
536665
const box = new THREE.Box3()
537666
let center = new THREE.Vector3()
@@ -629,13 +758,13 @@ app.registerExtension({
629758
load3d.remove()
630759
}
631760

632-
containerToLoad3D.delete(container)
761+
containerToLoad3D.delete(container.id)
633762

634-
origOnRemoved?.apply(this, arguments)
763+
origOnRemoved?.apply(this, [])
635764
}
636765

637766
node.onDrawBackground = function () {
638-
load3d.renderer.domElement.hidden = this.flags.collapsed
767+
load3d.renderer.domElement.hidden = this.flags.collapsed ?? false
639768
}
640769

641770
return {
@@ -689,6 +818,12 @@ app.registerExtension({
689818
const modelUrl = api.apiURL(getResourceURL(...splitFilePath(filename)))
690819

691820
load3d.loadModel(modelUrl, filename)
821+
822+
const material = node.widgets.find(
823+
(w: IWidget) => w.name === 'material'
824+
)
825+
826+
load3d.setMaterialMode(material.value)
692827
}
693828
}
694829

@@ -722,6 +857,14 @@ app.registerExtension({
722857
load3d.setViewPosition(value)
723858
}
724859

860+
const material = node.widgets.find((w: IWidget) => w.name === 'material')
861+
862+
material.callback = (value: 'original' | 'normal' | 'wireframe') => {
863+
load3d.setMaterialMode(value)
864+
}
865+
866+
load3d.setMaterialMode(material.value)
867+
725868
const bgColor = node.widgets.find((w: IWidget) => w.name === 'bg_color')
726869

727870
load3d.setBackgroundColor(bgColor.value)
@@ -774,7 +917,7 @@ app.registerExtension({
774917

775918
const fileInput = document.createElement('input')
776919
fileInput.type = 'file'
777-
fileInput.accept = '.gltf,.glb,.obj,.mtl,.fbx'
920+
fileInput.accept = '.gltf,.glb,.obj,.mtl,.fbx,.stl'
778921
fileInput.style.display = 'none'
779922
fileInput.onchange = () => {
780923
if (fileInput.files?.length) {

0 commit comments

Comments
 (0)