feat: Add UV reprojection and texture baking via WebGL

src/components/ModelViewer.tsx

@@ -7,6 +7,8 @@ import { useTexturePipeline } from "../hooks/useTexturePipeline";
 interface ModelViewerProps {
   object: THREE.Object3D | null;
   sceneObjects: any[];
+  bakeRequest?: { id: string, type: 'planar' | 'cylinder' | 'sphere' | 'box', timestamp: number } | null;
+  onBakeComplete?: (dataUrl: string) => void;
 }
 
 interface MeshMaterialManagerProps {
@@ -211,7 +213,9 @@ const FillMesh = ({ fillProps, originalMesh }: { fillProps: any, originalMesh: T
   return createPortal(fillNode, originalMesh);
 };
 
-export const ModelViewer = forwardRef<THREE.Group, ModelViewerProps>(({ object, sceneObjects }, ref) => {
+import { reprojectUVs, bakeTextureToNewUVs } from '../lib/bakingUtils';
+
+export const ModelViewer = forwardRef<THREE.Group, ModelViewerProps>(({ object, sceneObjects, bakeRequest, onBakeComplete }, ref) => {
   const internalRef = useRef<THREE.Group>(null);
   const groupRef = (ref as any) || internalRef;
   const textureCache = useRef<Map<string, THREE.Texture>>(new Map());
@@ -241,6 +245,43 @@ export const ModelViewer = forwardRef<THREE.Group, ModelViewerProps>(({ object,
     };
   }, []);
 
+  const { gl } = useThree();
+
+  // Handle Baking Request
+  useEffect(() => {
+    if (bakeRequest && onBakeComplete && meshCache.size > 0) {
+      const sceneObj = sceneObjects.find(obj => obj.id === bakeRequest.id);
+      if (!sceneObj) return;
+
+      // Find the primary mesh for this object (we assume the first mesh for now)
+      let targetMesh: THREE.Mesh | null = null;
+      for (const mesh of meshCache.values()) {
+        targetMesh = mesh;
+        break; 
+      }
+
+      if (targetMesh && sceneObj.materialProps?.map) {
+        const tex = textureCache.current.get(sceneObj.materialProps.map);
+        if (tex) {
+          // Keep old UVs
+          const oldUVArray = targetMesh.geometry.attributes.uv.array.slice() as Float32Array;
+          
+          // Apply new UV mapping (mutates geometry.attributes.uv)
+          reprojectUVs(targetMesh.geometry, bakeRequest.type);
+          
+          // Bake
+          try {
+            const dataUrl = bakeTextureToNewUVs(targetMesh, oldUVArray, tex, gl, 1024);
+            onBakeComplete(dataUrl);
+          } catch (e) {
+            console.error("Baking failed", e);
+            onBakeComplete(sceneObj.materialProps.map); // Fallback to original
+          }
+        }
+      }
+    }
+  }, [bakeRequest, meshCache, sceneObjects, gl, onBakeComplete]);
+
   // Main transform update for the core model
   useEffect(() => {
     if (object && sceneObjects) {

src/lib/bakingUtils.ts

@@ -0,0 +1,203 @@
+import * as THREE from 'three';
+
+/**
+ * Calculates new UVs for a geometry based on the chosen projection mapping.
+ * Supported types: 'planar', 'cylinder', 'sphere', 'box'
+ */
+export function reprojectUVs(geometry: THREE.BufferGeometry, mappingType: 'planar' | 'cylinder' | 'sphere' | 'box') {
+  if (!geometry.attributes.position) return;
+
+  geometry.computeBoundingBox();
+  const bbox = geometry.boundingBox!;
+  const size = new THREE.Vector3();
+  bbox.getSize(size);
+  const center = new THREE.Vector3();
+  bbox.getCenter(center);
+
+  const posAttr = geometry.attributes.position;
+  const uvArray = new Float32Array(posAttr.count * 2);
+
+  for (let i = 0; i < posAttr.count; i++) {
+    const x = posAttr.getX(i);
+    const y = posAttr.getY(i);
+    const z = posAttr.getZ(i);
+
+    // Normalize coordinates around center, scaled by size
+    const nx = size.x === 0 ? 0.5 : (x - bbox.min.x) / size.x;
+    const ny = size.y === 0 ? 0.5 : (y - bbox.min.y) / size.y;
+    const nz = size.z === 0 ? 0.5 : (z - bbox.min.z) / size.z;
+
+    let u = 0;
+    let v = 0;
+
+    if (mappingType === 'planar') {
+      // Front planar projection
+      u = nx;
+      v = ny;
+    } else if (mappingType === 'cylinder') {
+      // Cylindrical projection
+      const centeredX = x - center.x;
+      const centeredZ = z - center.z;
+      // atan2 returns -PI to PI. We want 0 to 1 around the cylinder.
+      let angle = Math.atan2(centeredZ, centeredX);
+      u = (angle + Math.PI) / (2 * Math.PI);
+      v = ny;
+    } else if (mappingType === 'sphere') {
+      // Spherical projection
+      const centeredX = x - center.x;
+      const centeredY = y - center.y;
+      const centeredZ = z - center.z;
+      const radius = Math.sqrt(centeredX * centeredX + centeredY * centeredY + centeredZ * centeredZ);
+      
+      const phi = Math.acos(centeredY / radius); // 0 to PI
+      let theta = Math.atan2(centeredZ, centeredX); // -PI to PI
+      
+      u = (theta + Math.PI) / (2 * Math.PI);
+      v = 1.0 - (phi / Math.PI);
+    } else if (mappingType === 'box') {
+      // Simple box projection (pick dominant axis based on normal)
+      // This requires normals
+      let normX = 0, normY = 0, normZ = 1;
+      if (geometry.attributes.normal) {
+        normX = geometry.attributes.normal.getX(i);
+        normY = geometry.attributes.normal.getY(i);
+        normZ = geometry.attributes.normal.getZ(i);
+      }
+      const ax = Math.abs(normX), ay = Math.abs(normY), az = Math.abs(normZ);
+      if (ax > ay && ax > az) {
+        u = nz; v = ny;
+      } else if (ay > ax && ay > az) {
+        u = nx; v = nz;
+      } else {
+        u = nx; v = ny;
+      }
+    }
+
+    // Clamp UVs to avoid edge artifacts
+    u = Math.max(0.001, Math.min(0.999, u));
+    v = Math.max(0.001, Math.min(0.999, v));
+
+    uvArray[i * 2] = u;
+    uvArray[i * 2 + 1] = v;
+  }
+
+  geometry.setAttribute('uv', new THREE.BufferAttribute(uvArray, 2));
+  geometry.attributes.uv.needsUpdate = true;
+}
+
+/**
+ * Bakes the old texture applied to the old UVs onto a new texture using the new UVs.
+ * We do this by rendering the mesh to a WebGLRenderTarget using a custom shader.
+ * @param mesh The mesh with the new UVs applied to geometry.attributes.uv
+ * @param oldUVArray The Float32Array of the original UVs
+ * @param oldTexture The original texture map
+ * @param renderer The active WebGLRenderer
+ * @param resolution The resolution of the baked texture
+ * @returns A data URL of the baked texture
+ */
+export function bakeTextureToNewUVs(
+  mesh: THREE.Mesh,
+  oldUVArray: Float32Array,
+  oldTexture: THREE.Texture,
+  renderer: THREE.WebGLRenderer,
+  resolution = 1024
+): string {
+  // Store original geometry attributes
+  const geometry = mesh.geometry.clone();
+  
+  // Create a secondary UV set for the old UVs
+  geometry.setAttribute('uv2', new THREE.BufferAttribute(oldUVArray, 2));
+
+  // Custom Shader:
+  // Position is derived from the new UVs (to flatten the mesh into the 2D canvas)
+  // Color is sampled from the old texture using the old UVs (uv2)
+  const bakeMaterial = new THREE.ShaderMaterial({
+    uniforms: {
+      tDiffuse: { value: oldTexture }
+    },
+    vertexShader: `
+      attribute vec2 uv2;
+      varying vec2 vOldUV;
+      void main() {
+        vOldUV = uv2;
+        // The new UVs (uv) range from 0 to 1.
+        // We map them to clip space: -1 to 1.
+        // Also flip Y so the texture is right-side up.
+        gl_Position = vec4(uv.x * 2.0 - 1.0, (uv.y * 2.0 - 1.0), 0.0, 1.0);
+      }
+    `,
+    fragmentShader: `
+      uniform sampler2D tDiffuse;
+      varying vec2 vOldUV;
+      void main() {
+        gl_FragColor = texture2D(tDiffuse, vOldUV);
+      }
+    `,
+    side: THREE.DoubleSide
+  });
+
+  const bakeMesh = new THREE.Mesh(geometry, bakeMaterial);
+  
+  const scene = new THREE.Scene();
+  scene.add(bakeMesh);
+  
+  // Orthographic camera covering exactly -1 to 1 clip space
+  const camera = new THREE.OrthographicCamera(-1, 1, 1, -1, -1, 1);
+  
+  // Render target for the bake
+  const renderTarget = new THREE.WebGLRenderTarget(resolution, resolution, {
+    format: THREE.RGBAFormat,
+    type: THREE.UnsignedByteType,
+    colorSpace: THREE.SRGBColorSpace,
+    magFilter: THREE.LinearFilter,
+    minFilter: THREE.LinearFilter
+  });
+
+  const originalRenderTarget = renderer.getRenderTarget();
+  const originalClearColor = renderer.getClearColor(new THREE.Color());
+  const originalClearAlpha = renderer.getClearAlpha();
+
+  // Render
+  renderer.setRenderTarget(renderTarget);
+  renderer.setClearColor(0x000000, 0); // Transparent background
+  renderer.clear();
+  renderer.render(scene, camera);
+
+  // Read pixels from render target
+  const pixels = new Uint8Array(resolution * resolution * 4);
+  renderer.readRenderTargetPixels(renderTarget, 0, 0, resolution, resolution, pixels);
+
+  // Restore renderer state
+  renderer.setRenderTarget(originalRenderTarget);
+  renderer.setClearColor(originalClearColor, originalClearAlpha);
+
+  // Flip Y when copying to canvas because readRenderTargetPixels reads bottom-up
+  const canvas = document.createElement('canvas');
+  canvas.width = resolution;
+  canvas.height = resolution;
+  const ctx = canvas.getContext('2d');
+  if (!ctx) throw new Error("Could not get 2d context");
+
+  const imgData = ctx.createImageData(resolution, resolution);
+  // Manual flip Y
+  for (let y = 0; y < resolution; y++) {
+    for (let x = 0; x < resolution; x++) {
+      const srcIndex = ((resolution - 1 - y) * resolution + x) * 4;
+      const destIndex = (y * resolution + x) * 4;
+      imgData.data[destIndex] = pixels[srcIndex];
+      imgData.data[destIndex + 1] = pixels[srcIndex + 1];
+      imgData.data[destIndex + 2] = pixels[srcIndex + 2];
+      imgData.data[destIndex + 3] = pixels[srcIndex + 3];
+    }
+  }
+
+  ctx.putImageData(imgData, 0, 0);
+
+  // Clean up
+  renderTarget.dispose();
+  geometry.dispose();
+  bakeMaterial.dispose();
+
+  // Return base64 JPEG to match sizes
+  return canvas.toDataURL('image/jpeg', 0.85);
+}

src/pages/Editor.tsx

@@ -638,6 +638,11 @@ export default function Editor() {
   } | null>(null);
   const [exportDialogName, setExportDialogName] = useState('');
   const [isRendering, setIsRendering] = useState(false);
+  const [rendering, setRendering] = useState(false);
+  const [renderProgress, setRenderProgress] = useState<{ step: string; progress: number } | null>(null);
+  const [isRenderComplete, setIsRenderComplete] = useState(false);
+  const [bakedDataUrl, setBakedDataUrl] = useState<string | null>(null);
+  const [bakeRequest, setBakeRequest] = useState<{ id: string, type: 'planar' | 'cylinder' | 'sphere' | 'box', timestamp: number } | null>(null);
   const [importStatus, setImportStatus] = useState<string | null>(null);
   const [importError, setImportError] = useState<string | null>(null);
   const [modelData, setModelData] = useState<{ name: string, data: string, extension: string } | null>(null);
@@ -2528,7 +2533,18 @@ export default function Editor() {
                 return null;
               })}
 
-              <ModelViewer ref={modelRef} object={loadedModel} sceneObjects={sceneObjects} />
+              <ModelViewer 
+                ref={modelRef} 
+                object={loadedModel} 
+                sceneObjects={sceneObjects} 
+                bakeRequest={bakeRequest}
+                onBakeComplete={(dataUrl) => {
+                  if (bakeRequest) {
+                    updateObjectMaterial(bakeRequest.id, { map: dataUrl, mappingType: 'uv' });
+                  }
+                  setBakeRequest(null);
+                }}
+              />
               
               {(!isRendering || !renderRequest?.includeAlpha) ? (
                 customHdri ? (
@@ -4667,8 +4683,24 @@ export default function Editor() {
                           <option value="planar">Planar</option>
                           <option value="box">Box</option>
                           <option value="cylinder">Cylinder</option>
+                          <option value="sphere">Sphere</option>
                           <option value="triplanar">Triplanar</option>
                         </select>
+                        {selectedObject.materialProps?.mappingType && selectedObject.materialProps.mappingType !== 'uv' && selectedObject.materialProps.mappingType !== 'triplanar' && (
+                          <button
+                            onClick={() => {
+                              if (!selectedObject.materialProps?.map) return;
+                              setBakeRequest({ 
+                                id: selectedObject.id, 
+                                type: selectedObject.materialProps.mappingType as any, 
+                                timestamp: Date.now() 
+                              });
+                            }}
+                            className="w-full py-2 bg-orange-500 hover:bg-orange-600 text-white font-bold rounded shadow transition-colors mt-1 text-[11px]"
+                          >
+                            {bakeRequest?.id === selectedObject.id ? 'Baking...' : 'Bake to UVs (Required for 2D AI)'}
+                          </button>
+                        )}
                       </div>
 
                       {/* UV Offset */}