lightningpixel · iammojogo-sudo · Jun 16, 2026
diff --git a/api/routers/optimize.py b/api/routers/optimize.py
@@ -12,6 +12,7 @@
     _pymeshlab = None
     _PYMESHLAB_AVAILABLE = False
 
+import numpy as np
 import trimesh
 import trimesh.visual
 from fastapi import APIRouter, HTTPException, UploadFile, File
@@ -35,6 +36,11 @@ class SmoothRequest(BaseModel):
     iterations: int
 
 
+class TransformRequest(BaseModel):
+    path: str                    # format: "{collection}/{filename}"
+    matrix: list[list[float]]    # row-major 4x4 world transform
+
+
 def _require_pymeshlab():
     if not _PYMESHLAB_AVAILABLE:
         raise HTTPException(503, "pymeshlab is unavailable on this system (DLL blocked by Windows Application Control policy)")
@@ -194,6 +200,34 @@ def smooth_mesh(body: SmoothRequest):
     return {"url": f"/workspace/{rel}"}
 
 
+@router.post("/transform")
+def transform_mesh(body: TransformRequest):
+    # Bake an interactive-gizmo transform into the GLB at scene level so it
+    # persists to export. Pure trimesh — no pymeshlab needed.
+    input_path = _resolve_input_path(body.path)
+
+    matrix = np.asarray(body.matrix, dtype=float)
+    if matrix.shape != (4, 4):
+        raise HTTPException(400, "matrix must be a 4x4 array")
+    if not np.all(np.isfinite(matrix)):
+        raise HTTPException(400, "matrix contains non-finite values")
+
+    # Keep the loaded result as-is (Scene when textured/multi-geometry) so
+    # apply_transform preserves materials and UVs.
+    loaded = trimesh.load(str(input_path))
+    loaded.apply_transform(matrix)
+
+    stem = input_path.stem
+    output_name = f"{stem}_xf_{uuid.uuid4().hex[:8]}.glb"
+    output_dir = input_path.parent if str(input_path).startswith(str(WORKSPACE_DIR.resolve())) else WORKSPACE_DIR / "Workflows"
+    output_dir.mkdir(parents=True, exist_ok=True)
+    output_path = output_dir / output_name
+    loaded.export(str(output_path))
+
+    rel = output_path.relative_to(WORKSPACE_DIR).as_posix()
+    return {"url": f"/workspace/{rel}"}
+
+
 def _smooth(input_path: str, iterations: int, tmp_dir: str) -> trimesh.Trimesh:
     loaded = trimesh.load(input_path)
     if isinstance(loaded, trimesh.Scene):
@@ -469,4 +503,4 @@ def export_mesh(path: str, format: str):
         content=data,
         media_type=mime,
         headers={"Content-Disposition": f'attachment; filename="{stem}.{format}"'},
-    )
+    )
diff --git a/src/areas/generate/GeneratePage.tsx b/src/areas/generate/GeneratePage.tsx
@@ -349,6 +349,22 @@ export default function GeneratePage(): JSX.Element {
     if (!meshSelected) setGizmoMode(null)
   }, [meshSelected])
 
+  // Gizmo hotkeys: W/E/R switch tool, Esc exits. Ignored while typing.
+  useEffect(() => {
+    const handler = (e: KeyboardEvent) => {
+      const el = document.activeElement as HTMLElement | null
+      if (el && (el instanceof HTMLInputElement || el instanceof HTMLTextAreaElement || el.isContentEditable)) return
+      if (e.key === 'Escape') { setGizmoMode((m) => (m ? null : m)); return }
+      if (!hasModel || !meshSelected) return
+      const k = e.key.toLowerCase()
+      if (k === 'w') setGizmoMode('translate')
+      else if (k === 'e') setGizmoMode('rotate')
+      else if (k === 'r') setGizmoMode('scale')
+    }
+    window.addEventListener('keydown', handler)
+    return () => window.removeEventListener('keydown', handler)
+  }, [hasModel, meshSelected])
+
   async function handleUnloadAll() {
     await window.electron.model.unloadAll()
     setUnloadStatus('done')
@@ -736,10 +752,10 @@ export default function GeneratePage(): JSX.Element {
 
         {/* Viewer area */}
         <div className="flex-1 relative overflow-hidden">
-          <Viewer3D lightSettings={lightSettings} />
+          <Viewer3D lightSettings={lightSettings} gizmoMode={gizmoMode} />
           <GenerationHUD />
         </div>
       </div>
     </>
   )
-}
+}
diff --git a/src/areas/generate/components/Viewer3D.tsx b/src/areas/generate/components/Viewer3D.tsx
@@ -1,7 +1,7 @@
 import { Component, Suspense, useEffect, useMemo, useRef, useState } from 'react'
 import type { ReactNode, ErrorInfo } from 'react'
 import { Canvas, useLoader, useThree } from '@react-three/fiber'
-import { Environment, GizmoHelper, Lightformer, OrbitControls, useGizmoContext, useGLTF } from '@react-three/drei'
+import { Environment, GizmoHelper, Lightformer, OrbitControls, TransformControls, useGizmoContext, useGLTF } from '@react-three/drei'
 import { EffectComposer, Outline, Select, Selection } from '@react-three/postprocessing'
 import * as THREE from 'three'
 import { OBJLoader } from 'three/examples/jsm/loaders/OBJLoader.js'
@@ -13,11 +13,14 @@ THREE.BufferGeometry.prototype.disposeBoundsTree = disposeBoundsTree as any
 THREE.Mesh.prototype.raycast = acceleratedRaycast
 import SplatViewer, { type SplatViewerHandle } from './SplatViewer'
 import { useGeneration } from '@shared/hooks/useGeneration'
+import { useApi } from '@shared/hooks/useApi'
 import { useAppStore } from '@shared/stores/appStore'
 import { ViewerToolbar, type ViewMode } from './ViewerToolbar'
 import type { LightSettings } from '../GeneratePage'
 import { DEFAULT_LIGHT_SETTINGS } from '../GeneratePage'
 
+export type GizmoMode = 'translate' | 'rotate' | 'scale'
+
 const SELECTION_OUTLINE_VISIBLE_COLOR = 0x8b5cf6
 const SELECTION_OUTLINE_HIDDEN_COLOR = 0x5b21b6
 const SELECTION_OUTLINE_EDGE_STRENGTH = 2.5
@@ -135,13 +138,16 @@ interface MeshModelProps {
   jobId: string
   viewMode: ViewMode
   selected: boolean
+  autoCenter: boolean
+  resetToken: number
   onStats: (stats: { vertices: number; triangles: number }) => void
   onSelect: () => void
+  onObject: (obj: THREE.Object3D | null) => void
 }
 
-function MeshModel({ url, jobId, viewMode, selected, onStats, onSelect }: MeshModelProps): JSX.Element {
+function MeshModel({ url, jobId, viewMode, selected, autoCenter, resetToken, onStats, onSelect, onObject }: MeshModelProps): JSX.Element {
   const extension = url.split('?')[0]?.split('.').pop()?.toLowerCase()
-  const common = { url, jobId, viewMode, selected, onStats, onSelect }
+  const common = { url, jobId, viewMode, selected, autoCenter, resetToken, onStats, onSelect, onObject }
   return extension === 'obj' ? <ObjMeshModel {...common} /> : <GltfMeshModel {...common} />
 }
 
@@ -159,8 +165,11 @@ function SceneMeshModel({
   url,
   viewMode,
   selected,
+  autoCenter,
+  resetToken,
   onStats,
   onSelect,
+  onObject,
   scene,
   loaderType,
 }: MeshModelProps & {
@@ -170,6 +179,12 @@ function SceneMeshModel({
   const captured = useRef(false)
   const edgeHelpers = useRef<THREE.LineSegments[]>([])
 
+  // Expose the scene object so Viewer3D can attach the transform gizmo to it.
+  useEffect(() => {
+    onObject(scene)
+    return () => onObject(null)
+  }, [scene, onObject])
+
   // Free GPU resources and loader cache when this model is replaced or unmounted
   useEffect(() => {
     return () => {
@@ -208,17 +223,23 @@ function SceneMeshModel({
     }
   }, [scene])
 
-  // Centre the mesh on the grid
+  // Centre the mesh on the grid.
+  // Only runs on first load / model change (autoCenter) or an explicit Reset
+  // (resetToken) — never on plain re-renders, so a live gizmo transform or a
+  // baked "Apply" pose is not silently overwritten.
   useEffect(() => {
-    // Reset before computing — useGLTF caches the scene with its already-modified position,
-    // which would skew the setFromObject (world space) on a second mount.
-    scene.position.set(0, 0, 0)
-    const box = new THREE.Box3().setFromObject(scene)
-    const center = new THREE.Vector3()
-    box.getCenter(center)
-    scene.position.set(-center.x, -box.min.y, -center.z)
-
-    // Compute stats
+    if (autoCenter) {
+      // Clear any live gizmo transform before measuring.
+      scene.position.set(0, 0, 0)
+      scene.rotation.set(0, 0, 0)
+      scene.scale.set(1, 1, 1)
+      const box = new THREE.Box3().setFromObject(scene)
+      const center = new THREE.Vector3()
+      box.getCenter(center)
+      scene.position.set(-center.x, -box.min.y, -center.z)
+    }
+
+    // Compute stats (independent of centering)
     let vertices = 0
     let triangles = 0
     scene.traverse((child) => {
@@ -231,7 +252,7 @@ function SceneMeshModel({
     })
     const roundedTriangles = Math.round(triangles)
     onStats({ vertices: Math.round(vertices), triangles: roundedTriangles })
-  }, [scene])
+  }, [scene, autoCenter, resetToken])
 
   // Thumbnail capture (kept for future use)
   useEffect(() => {
@@ -380,13 +401,17 @@ function EmptyState(): JSX.Element {
 // Viewer3D
 // ---------------------------------------------------------------------------
 
-export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS }: { lightSettings?: LightSettings }): JSX.Element {
+export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS, gizmoMode = null }: { lightSettings?: LightSettings; gizmoMode?: GizmoMode | null }): JSX.Element {
   const { currentJob } = useGeneration()
   const apiUrl = useAppStore((s) => s.apiUrl)
 
   const setStoreMeshStats = useAppStore((s) => s.setMeshStats)
   const meshStats = useAppStore((s) => s.meshStats)
   const setCurrentJob = useAppStore((s) => s.setCurrentJob)
+  const updateCurrentJob = useAppStore((s) => s.updateCurrentJob)
+  const pushMeshUrl = useAppStore((s) => s.pushMeshUrl)
+  const showError = useAppStore((s) => s.showError)
+  const { transformMesh } = useApi()
 
   const [viewMode, setViewMode] = useState<ViewMode>('solid')
   const [autoRotate, setAutoRotate] = useState(false)
@@ -395,12 +420,21 @@ export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS }: { l
   const canvasRef = useRef<HTMLCanvasElement | null>(null)
   const splatRef = useRef<SplatViewerHandle | null>(null)
 
+  const [meshObject, setMeshObject] = useState<THREE.Object3D | null>(null)
+  const [resetToken, setResetToken] = useState(0)
+  const [applying, setApplying] = useState(false)
+  // URLs whose geometry already has a baked transform — these must NOT be
+  // re-centered on load, so the applied pose is shown verbatim.
+  const appliedUrls = useRef<Set<string>>(new Set())
+
   const outputUrl = currentJob?.outputUrl ?? ''
   const modelUrl =
     currentJob?.status === 'done' && currentJob.outputUrl
       ? `${apiUrl}${currentJob.outputUrl}`
       : null
 
+  const autoCenter = !appliedUrls.current.has(outputUrl)
+
   // A .ply/.splat reaching the viewer is always a Gaussian splat here: mesh
   // plys are converted to GLB on import and workflow mesh outputs are .glb.
   const isSplat = /\.(ply|splat)$/i.test(outputUrl)
@@ -446,6 +480,50 @@ export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS }: { l
     link.click()
   }
 
+  const handleResetTransform = () => {
+    const original = currentJob?.originalOutputUrl
+    if (original && currentJob?.outputUrl !== original) {
+      // Was baked via Apply — reload the original (auto-centred on load).
+      updateCurrentJob({ outputUrl: original })
+      pushMeshUrl(original)
+    } else {
+      // Live transform only — re-centre the current scene in place.
+      setResetToken((t) => t + 1)
+    }
+  }
+
+  const handleApplyTransform = async () => {
+    if (!meshObject || !currentJob?.outputUrl) return
+    const url = currentJob.outputUrl
+    if (!url.startsWith('/workspace/')) return
+    const path = url.slice('/workspace/'.length)
+
+    // We bake the full world matrix, which includes the centering offset the
+    // viewer applies on load ("bake what you see"). The result URL is then
+    // flagged so it is NOT re-centered on reload, keeping the visible pose.
+    meshObject.updateWorldMatrix(true, false)
+    const e = meshObject.matrixWorld.elements
+    // THREE stores column-major; emit a row-major 4x4 for the backend.
+    const matrix = [
+      [e[0], e[4], e[8], e[12]],
+      [e[1], e[5], e[9], e[13]],
+      [e[2], e[6], e[10], e[14]],
+      [e[3], e[7], e[11], e[15]],
+    ]
+
+    setApplying(true)
+    try {
+      const result = await transformMesh(path, matrix)
+      appliedUrls.current.add(result.url)
+      updateCurrentJob({ outputUrl: result.url })
+      pushMeshUrl(result.url)
+    } catch (err) {
+      showError(err instanceof Error ? err.message : String(err))
+    } finally {
+      setApplying(false)
+    }
+  }
+
 
   return (
     <ModelErrorBoundary resetKey={modelUrl} fallback={<ModelLoadError />}>
@@ -504,13 +582,20 @@ export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS }: { l
                   jobId={currentJob.id}
                   viewMode={viewMode}
                   selected={selected}
+                  autoCenter={autoCenter}
+                  resetToken={resetToken}
                   onStats={setStoreMeshStats}
                   onSelect={() => setSelected(true)}
+                  onObject={setMeshObject}
                 />
               </Suspense>
             </Selection>
           ) : null}
 
+          {selected && gizmoMode && meshObject && (
+            <TransformControls object={meshObject} mode={gizmoMode} />
+          )}
+
           <OrbitControls
             makeDefault
             enablePan
@@ -542,6 +627,33 @@ export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS }: { l
           />
         )}
 
+        {/* Transform apply/reset — visible while a gizmo tool is active */}
+        {!isSplat && modelUrl && selected && gizmoMode && (
+          <div className="absolute top-4 left-1/2 -translate-x-1/2 z-20 flex items-center gap-1.5 bg-zinc-900/80 border border-zinc-700/60 backdrop-blur-sm rounded-lg px-1.5 py-1">
+            <button
+              onClick={handleApplyTransform}
+              disabled={applying}
+              className="flex items-center gap-1.5 px-2.5 h-7 rounded-md text-xs font-medium bg-violet-600 text-white hover:bg-violet-500 disabled:opacity-40 disabled:pointer-events-none transition-colors"
+            >
+              <svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2">
+                <path d="M20 6L9 17l-5-5" />
+              </svg>
+              {applying ? 'Applying\u2026' : 'Apply'}
+            </button>
+            <button
+              onClick={handleResetTransform}
+              disabled={applying}
+              className="flex items-center gap-1.5 px-2.5 h-7 rounded-md text-xs font-medium text-zinc-300 hover:text-white hover:bg-zinc-700/60 disabled:opacity-40 disabled:pointer-events-none transition-colors"
+            >
+              <svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" strokeWidth="2">
+                <path d="M3 7v6h6" />
+                <path d="M21 17A9 9 0 0 0 6 10.3L3 13" />
+              </svg>
+              Reset
+            </button>
+          </div>
+        )}
+
         {/* Bottom-left stats overlay */}
         {meshStats && (
           <div className="absolute bottom-4 left-4 pointer-events-none">
@@ -565,4 +677,4 @@ export default function Viewer3D({ lightSettings = DEFAULT_LIGHT_SETTINGS }: { l
       </div>
     </ModelErrorBoundary>
   )
-}
+}
diff --git a/src/shared/hooks/useApi.ts b/src/shared/hooks/useApi.ts
@@ -116,5 +116,18 @@ export function useApi() {
     return { url: data.url }
   }
 
-  return { generateFromImage, pollJobStatus, cancelJob, getModelStatus, downloadModel, optimizeMesh, smoothMesh, importMesh }
-}
+  // Bakes a world-space 4x4 transform into the GLB geometry.
+  // `matrix` is row-major (4 rows of 4), matching the backend's reshape.
+  async function transformMesh(
+    path: string,
+    matrix: number[][],
+  ): Promise<{ url: string }> {
+    const { data } = await client.post<{ url: string }>('/optimize/transform', {
+      path,
+      matrix,
+    })
+    return { url: data.url }
+  }
+
+  return { generateFromImage, pollJobStatus, cancelJob, getModelStatus, downloadModel, optimizeMesh, smoothMesh, importMesh, transformMesh }
+}