refactor screenshot code

applications/visualizer/frontend/src/components/viewers/ThreeD/SceneControls.tsx

@@ -7,11 +7,10 @@ import { useRef, useState } from "react";
 import { vars } from "../../../theme/variables.ts";
 import CustomFormControlLabel from "./CustomFormControlLabel.tsx";
 import { Recorder } from "./Recorder.ts";
-import { downloadScreenshot } from "./Screenshoter.ts";
 
 const { gray500 } = vars;
 
-function SceneControls({ cameraControlRef, isWireframe, setIsWireframe, recorderRef }) {
+function SceneControls({ cameraControlRef, isWireframe, setIsWireframe, recorderRef, handleScreenshot }) {
   const [anchorEl, setAnchorEl] = useState<HTMLElement | null>(null);
   const rotateAnimationRef = useRef<number | null>(null);
   const [isRotating, setIsRotating] = useState(false);
@@ -56,13 +55,6 @@ function SceneControls({ cameraControlRef, isWireframe, setIsWireframe, recorder
 
     setIsRotating(!isRotating);
   };
-
-  const handleScreenshot = () => {
-    if (cameraControlRef.current) {
-      downloadScreenshot(document.getElementsByTagName("canvas")[0], 0.95, { width: 3840, height: 2160 }, 1, () => true, "screenshot.png");
-    }
-  };
-
   const startRecording = () => {
     if (recorderRef.current === null) {
       const canvas = document.getElementsByTagName("canvas")[0];

applications/visualizer/frontend/src/components/viewers/ThreeD/Screenshoter.ts

@@ -1,36 +1,6 @@
-import * as htmlToImage from "html-to-image";
-import { formatDate } from "../../../helpers/utils.ts";
+import * as THREE from "three";
 
-function getOptions(htmlElement, targetResolution, quality, pixelRatio, filter) {
-  const resolution = getResolutionFixedRatio(htmlElement, targetResolution);
-  return {
-    quality: quality,
-    canvasWidth: resolution.width,
-    canvasHeight: resolution.height,
-    pixelRatio: pixelRatio,
-    filter: filter,
-  };
-}
-
-export function downloadScreenshot(
-  htmlElement,
-  quality = 0.95,
-  targetResolution = { width: 3840, height: 2160 },
-  pixelRatio = 1,
-  filter = () => true,
-  filename = `Canvas_${formatDate(new Date())}.png`,
-) {
-  const options = getOptions(htmlElement, targetResolution, quality, pixelRatio, filter);
-
-  htmlToImage.toBlob(htmlElement, options).then((blob) => {
-    const link = document.createElement("a");
-    link.download = filename;
-    link.href = window.URL.createObjectURL(blob);
-    link.click();
-  });
-}
-
-function getResolutionFixedRatio(htmlElement, target) {
+function getResolutionFixedRatio(htmlElement: HTMLElement, target: { width: number; height: number }) {
   const current = {
     height: htmlElement.clientHeight,
     width: htmlElement.clientWidth,
@@ -47,3 +17,47 @@ function getResolutionFixedRatio(htmlElement, target) {
     width: target.width,
   };
 }
+
+function getOptions(htmlElement: HTMLCanvasElement, targetResolution: { width: number; height: number }, pixelRatio: number) {
+  const resolution = getResolutionFixedRatio(htmlElement, targetResolution);
+  return {
+    canvasWidth: resolution.width,
+    canvasHeight: resolution.height,
+    pixelRatio: pixelRatio,
+  };
+}
+
+export function downloadScreenshot(
+  canvasRef: React.RefObject<HTMLCanvasElement>,
+  sceneRef: React.RefObject<THREE.Scene>,
+  cameraRef: React.RefObject<THREE.PerspectiveCamera>,
+) {
+  if (!sceneRef.current || !cameraRef.current || !canvasRef.current) return;
+
+  const options = getOptions(canvasRef.current, { width: 3840, height: 2160 }, 1);
+
+  try {
+    const tempRenderer = new THREE.WebGLRenderer({ preserveDrawingBuffer: true });
+    tempRenderer.setSize(options.canvasWidth, options.canvasHeight);
+    tempRenderer.setPixelRatio(options.pixelRatio); // Set the resolution scaling
+
+    cameraRef.current.aspect = options.canvasWidth / options.canvasHeight;
+    cameraRef.current.updateProjectionMatrix();
+
+    tempRenderer.render(sceneRef.current, cameraRef.current);
+
+    tempRenderer.domElement.toBlob((blob) => {
+      if (blob) {
+        const link = document.createElement("a");
+        link.href = URL.createObjectURL(blob);
+        link.download = "screenshot.png";
+        link.click();
+        URL.revokeObjectURL(link.href);
+      }
+    }, "image/png");
+
+    tempRenderer.dispose();
+  } catch (e) {
+    console.error("Error saving image:", e);
+  }
+}

applications/visualizer/frontend/src/components/viewers/ThreeD/ThreeDViewer.tsx

@@ -1,6 +1,7 @@
 import { CameraControls, PerspectiveCamera } from "@react-three/drei";
 import { Canvas } from "@react-three/fiber";
 import { Suspense, useEffect, useMemo, useRef, useState } from "react";
+import type * as THREE from "three";
 import { useSelectedWorkspace } from "../../../hooks/useSelectedWorkspace.ts";
 import { ViewerType, getNeuronUrlForDataset } from "../../../models/models.ts";
 import { type Dataset, OpenAPI } from "../../../rest";
@@ -20,7 +21,7 @@ import Loader from "./Loader.tsx";
 import type { Recorder } from "./Recorder";
 import STLViewer from "./STLViewer.tsx";
 import SceneControls from "./SceneControls.tsx";
-
+import { downloadScreenshot } from "./Screenshoter.ts";
 export interface Instance {
   id: string;
   url: string;
@@ -37,8 +38,11 @@ function ThreeDViewer() {
   const [isWireframe, setIsWireframe] = useState<boolean>(false);
 
   const cameraControlRef = useRef<CameraControls | null>(null);
-
   const recorderRef = useRef<Recorder | null>(null);
+  const canvasRef = useRef<HTMLCanvasElement | null>(null);
+  const sceneRef = useRef<THREE.Scene | null>(null);
+  const cameraRef = useRef<THREE.PerspectiveCamera | null>(null);
+  const glRef = useRef<THREE.WebGLRenderer | null>(null);
 
   // @ts-expect-error 'setShowNeurons' is declared but its value is never read.
   // eslint-disable-next-line @typescript-eslint/no-unused-vars
@@ -67,10 +71,21 @@ function ThreeDViewer() {
     setInstances(newInstances);
   }, [selectedDataset, workspace.availableNeurons, workspace.visibilities]);
 
+  const handleScreenshot = () => {
+    downloadScreenshot(canvasRef, sceneRef, cameraRef);
+  };
+
+  const onCreated = (state) => {
+    canvasRef.current = state.gl.domElement;
+    sceneRef.current = state.scene;
+    cameraRef.current = state.camera;
+    glRef.current = state.gl;
+  };
+
   return (
     <>
       <DatasetPicker datasets={dataSets} selectedDataset={selectedDataset} onDatasetChange={setSelectedDataset} />
-      <Canvas style={{ backgroundColor: SCENE_BACKGROUND }} frameloop={"demand"} gl={{ preserveDrawingBuffer: true }}>
+      <Canvas style={{ backgroundColor: SCENE_BACKGROUND }} frameloop={"demand"} gl={{ preserveDrawingBuffer: false }} onCreated={onCreated}>
         <color attach="background" args={["#F6F5F4"]} />
         <Suspense fallback={<Loader />}>
           <PerspectiveCamera
@@ -80,6 +95,7 @@ function ThreeDViewer() {
             position={CAMERA_POSITION}
             near={CAMERA_NEAR}
             far={CAMERA_FAR}
+            ref={cameraRef}
           />
           <CameraControls ref={cameraControlRef} />
 
@@ -91,7 +107,13 @@ function ThreeDViewer() {
           <STLViewer instances={instances} isWireframe={isWireframe} />
         </Suspense>
       </Canvas>
-      <SceneControls cameraControlRef={cameraControlRef} isWireframe={isWireframe} setIsWireframe={setIsWireframe} recorderRef={recorderRef} />
+      <SceneControls
+        cameraControlRef={cameraControlRef}
+        isWireframe={isWireframe}
+        setIsWireframe={setIsWireframe}
+        recorderRef={recorderRef}
+        handleScreenshot={handleScreenshot}
+      />
     </>
   );
 }