From 43d0d443fe65b99734cfee055be4aad695d9e2f4 Mon Sep 17 00:00:00 2001
From: bflatastic <bombayapirat@gmail.com>
Date: Mon, 28 Jul 2025 13:20:26 +0700
Subject: [PATCH] snapshot bug fix

---
 .../app/src/pages/PageDraw/PageDraw.tsx       |   9 +-
 .../pages/PageDraw/hooks/useUpdateSnapshot.ts | 211 ++++++++++--------
 2 files changed, 121 insertions(+), 99 deletions(-)

diff --git a/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx b/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx
index 40414350a9..79c6371d5b 100644
--- a/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx
+++ b/frontend/apps/editor3d/app/src/pages/PageDraw/PageDraw.tsx
@@ -170,8 +170,15 @@ const PageDraw = () => {
   };
 
   const onEnqueuedPressed = async () => {
+    const { width, height } = store.getAspectRatioDimensions();
+
     // takes snap shot and then a global variable in the engine will invoke the inference.
-    const image = await captureStageImageBitmap(stageRef, transformerRefs);
+    const image = await captureStageImageBitmap(
+      stageRef,
+      transformerRefs,
+      width,
+      height,
+    );
     if (!image) {
       console.error("Failed to capture stage image");
       return;
diff --git a/frontend/apps/editor3d/app/src/pages/PageDraw/hooks/useUpdateSnapshot.ts b/frontend/apps/editor3d/app/src/pages/PageDraw/hooks/useUpdateSnapshot.ts
index 8aa7f86d18..4bf2a868e8 100644
--- a/frontend/apps/editor3d/app/src/pages/PageDraw/hooks/useUpdateSnapshot.ts
+++ b/frontend/apps/editor3d/app/src/pages/PageDraw/hooks/useUpdateSnapshot.ts
@@ -1,114 +1,129 @@
-import { useState, useEffect } from 'react';
+import { useState, useEffect } from "react";
+import Konva from "konva";
 
-const debugOutput = (uri:string, name:string)=> {
-    var link = document.createElement('a');
-    link.download = name;
-    link.href = uri;
-    document.body.appendChild(link);
-    link.click();
-    document.body.removeChild(link);
+export const captureStageImageBitmap = async (
+  stageRef: React.RefObject<Konva.Stage>,
+  transformerRefs: React.RefObject<{ [key: string]: Konva.Transformer }>,
+  width: number = 1024,
+  height: number = 1024,
+): Promise<ImageBitmap | undefined> => {
+  if (!stageRef.current) {
+    return undefined;
   }
 
+  const oldPos = { ...stageRef.current.position() };
+  const oldScale = {
+    x: stageRef.current.scaleX(),
+    y: stageRef.current.scaleY(),
+  };
 
- export const captureStageImageBitmap = async (stageRef: any,transformerRefs: any): Promise<ImageBitmap | undefined> => {
-    if (!stageRef.current) {
-      return undefined;
-    }
-  
-    const oldPos = { ...stageRef.current.position() };
-    const oldScale = { x: stageRef.current.scaleX(), y: stageRef.current.scaleY() };
-  
-    // Hide all transformers
-    Object.values(transformerRefs.current).forEach(transformer => {
+  // Hide all transformers
+  Object.values(transformerRefs.current || {}).forEach(
+    (transformer: Konva.Transformer) => {
       transformer.visible(false);
+    },
+  );
+
+  stageRef.current.position({ x: 0, y: 0 });
+  stageRef.current.scale({ x: 1, y: 1 });
+  stageRef.current.draw();
+
+  try {
+    const canvas = stageRef.current.toCanvas({
+      width,
+      height,
+      pixelRatio: 1,
     });
-  
+
+    return await createImageBitmap(canvas);
+  } catch (error) {
+    console.error("Failed to create ImageBitmap:", error);
+    return undefined;
+  } finally {
+    // Restore transformers
+    Object.values(transformerRefs.current || {}).forEach(
+      (transformer: Konva.Transformer) => {
+        transformer.visible(true);
+      },
+    );
+
+    stageRef.current.position(oldPos);
+    stageRef.current.scale(oldScale);
+    stageRef.current.draw();
+  }
+};
+
+export const useStageSnapshot = (
+  stageRef: React.RefObject<Konva.Stage>,
+  imageRef: React.RefObject<Konva.Image>,
+  isSelectingRef: React.RefObject<boolean>,
+  transformerRefs: React.RefObject<{ [key: string]: Konva.Transformer }>,
+) => {
+  const [animationFrameId, setAnimationFrameId] = useState<number | null>(null);
+
+  const captureStageImage = (
+    stageRef: React.RefObject<Konva.Stage>,
+    width: number = 1024,
+    height: number = 1024,
+  ): string => {
+    if (!stageRef.current) return "";
+
+    const oldPos = { ...stageRef.current.position() };
+    const oldScale = {
+      x: stageRef.current.scaleX(),
+      y: stageRef.current.scaleY(),
+    };
+
+    // Hide all transformers
+    Object.values(transformerRefs.current || {}).forEach(
+      (transformer: Konva.Transformer) => {
+        transformer.visible(false);
+      },
+    );
+
     stageRef.current.position({ x: 0, y: 0 });
     stageRef.current.scale({ x: 1, y: 1 });
     stageRef.current.draw();
-  
-    try {
-      const canvas = stageRef.current.toCanvas({ 
-        width: 1024, 
-        height: 1024, 
-        pixelRatio: 1 
-      });
-      
-      return await createImageBitmap(canvas);
-    } catch (error) {
-      console.error('Failed to create ImageBitmap:', error);
-      return undefined;
-    } finally {
-      // Restore transformers
-      Object.values(transformerRefs.current).forEach(transformer => {
+    const uri = stageRef.current.toDataURL({ width, height, pixelRatio: 1 });
+
+    // Restore transformers
+    Object.values(transformerRefs.current || {}).forEach(
+      (transformer: Konva.Transformer) => {
         transformer.visible(true);
-      });
-  
-      stageRef.current.position(oldPos);
-      stageRef.current.scale(oldScale);
-      stageRef.current.draw();
-    }
+      },
+    );
+
+    stageRef.current.position(oldPos);
+    stageRef.current.scale(oldScale);
+    stageRef.current.draw();
+    return uri;
   };
 
- export const useStageSnapshot = (
-    stageRef: React.RefObject<any>,
-    imageRef: React.RefObject<any>,
-    isSelectingRef: React.RefObject<boolean>,
-    transformerRefs: React.RefObject<{ [key: string]: Konva.Transformer }>
-  ) => {
-    const [animationFrameId, setAnimationFrameId] = useState<number | null>(null);
-
-    const captureStageImage = (stageRef: any): string => {
-      const oldPos = { ...stageRef.current.position() };
-      const oldScale = { x: stageRef.current.scaleX(), y: stageRef.current.scaleY() };
-    
-      // Hide all transformers
-      Object.values(transformerRefs.current).forEach(transformer => {
-        transformer.visible(false);
-      });
-    
-      stageRef.current.position({ x: 0, y: 0 });
-      stageRef.current.scale({ x: 1, y: 1 });
-      stageRef.current.draw();
-      const uri = stageRef.current.toDataURL({ width: 1024, height: 1024, pixelRatio: 1 });
-  
-      // Restore transformers
-      Object.values(transformerRefs.current).forEach(transformer => {
-        transformer.visible(true);
-      });
-    
-      stageRef.current.position(oldPos);
-      stageRef.current.scale(oldScale);
-      stageRef.current.draw();
-      return uri;
-    };
-
-    useEffect(() => {
-      const updateSnapshot = () => {
-        if (!isSelectingRef.current && stageRef.current && imageRef.current) {
-          const uri = captureStageImage(stageRef);
-          if (uri) {
-            const img = new window.Image();
-            img.onload = () => {
-              imageRef.current?.setAttrs({
-                image: img
-              });
-            };
-            img.src = uri;
-          }
+  useEffect(() => {
+    const updateSnapshot = () => {
+      if (!isSelectingRef.current && stageRef.current && imageRef.current) {
+        const uri = captureStageImage(stageRef);
+        if (uri) {
+          const img = new window.Image();
+          img.onload = () => {
+            imageRef.current?.setAttrs({
+              image: img,
+            });
+          };
+          img.src = uri;
         }
-        const frameId = requestAnimationFrame(updateSnapshot);
-        setAnimationFrameId(frameId);
-      };
-    
+      }
       const frameId = requestAnimationFrame(updateSnapshot);
       setAnimationFrameId(frameId);
-    
-      return () => {
-        if (animationFrameId) {
-          cancelAnimationFrame(animationFrameId);
-        }
+    };
 
-      };
-    }, []);
-  };
\ No newline at end of file
+    const frameId = requestAnimationFrame(updateSnapshot);
+    setAnimationFrameId(frameId);
+
+    return () => {
+      if (animationFrameId) {
+        cancelAnimationFrame(animationFrameId);
+      }
+    };
+  }, []);
+};