refactor screenshot code

applications/visualizer/frontend/package.json

@@ -31,6 +31,7 @@
     "cytoscape-dagre": "^2.5.0",
     "cytoscape-fcose": "^2.2.0",
     "file-saver": "^2.0.5",
+    "html-to-image": "^1.11.11",
     "html2canvas": "^1.4.1",
     "ol": "^9.1.0",
     "pako": "^2.1.0",

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

@@ -0,0 +1,115 @@
+import { formatDate } from "../../../helpers/utils.ts";
+export class Recorder {
+  private mediaRecorder: MediaRecorder | null = null;
+  private recordedBlobs: Blob[] = [];
+  private stream: MediaStream;
+  private ctx: WebGLRenderingContext;
+  // @ts-ignore
+  private options: { mediaRecorderOptions?: MediaRecorderOptions; blobOptions?: BlobPropertyBag } = {
+    mediaRecorderOptions: { mimeType: "video/webm" },
+    blobOptions: { type: "video/webm" },
+  };
+  private blobOptions: BlobPropertyBag = { type: "video/webm" };
+
+  constructor(canvas: HTMLCanvasElement, recorderOptions: { mediaRecorderOptions?: MediaRecorderOptions; blobOptions?: BlobPropertyBag }) {
+    this.stream = canvas.captureStream();
+    const { mediaRecorderOptions, blobOptions } = recorderOptions;
+    this.setupMediaRecorder(mediaRecorderOptions);
+    this.recordedBlobs = [];
+    this.blobOptions = blobOptions;
+    this.ctx = canvas.getContext("webgl");
+  }
+
+  handleDataAvailable(event) {
+    if (event.data && event.data.size > 0) {
+      this.recordedBlobs.push(event.data);
+    }
+  }
+
+  setupMediaRecorder(options) {
+    let error = "";
+
+    if (options == null) {
+      options = { mimeType: "video/webm" };
+    }
+    let mediaRecorder;
+    try {
+      mediaRecorder = new MediaRecorder(this.stream, options);
+    } catch (e0) {
+      error = `Unable to create MediaRecorder with options Object: ${e0}`;
+      try {
+        options = { mimeType: "video/webm,codecs=vp9" };
+        mediaRecorder = new MediaRecorder(this.stream, options);
+      } catch (e1) {
+        error = `Unable to create MediaRecorder with options Object: ${e1}`;
+        try {
+          options = { mimeType: "video/webm,codecs=vp8" }; // Chrome 47
+          mediaRecorder = new MediaRecorder(this.stream, options);
+        } catch (e2) {
+          error =
+            "MediaRecorder is not supported by this browser.\n\n" +
+            "Try Firefox 29 or later, or Chrome 47 or later, " +
+            "with Enable experimental Web Platform features enabled from chrome://flags." +
+            `Exception while creating MediaRecorder: ${e2}`;
+        }
+      }
+    }
+
+    if (!mediaRecorder) {
+      throw new Error(error);
+    }
+
+    mediaRecorder.ondataavailable = (evt) => this.handleDataAvailable(evt);
+    mediaRecorder.onstart = () => this.animationLoop();
+
+    this.mediaRecorder = mediaRecorder;
+    this.options = options;
+    if (!this.blobOptions) {
+      const { mimeType } = options;
+      this.blobOptions = { type: mimeType };
+    }
+  }
+
+  startRecording() {
+    this.recordedBlobs = [];
+    this.mediaRecorder.start(100);
+  }
+
+  stopRecording(options) {
+    this.mediaRecorder.stop();
+    return this.getRecordingBlob(options);
+  }
+
+  download(filename, options) {
+    if (!filename) {
+      filename = `CanvasRecording_${formatDate(new Date())}.webm`;
+    }
+    const blob = this.getRecordingBlob(options);
+    const url = window.URL.createObjectURL(blob);
+    const a = document.createElement("a");
+    a.style.display = "none";
+    a.href = url;
+    a.download = filename;
+    document.body.appendChild(a);
+    a.click();
+    setTimeout(() => {
+      document.body.removeChild(a);
+      window.URL.revokeObjectURL(url);
+    }, 100);
+    return blob;
+  }
+
+  getRecordingBlob(options) {
+    if (!options) {
+      options = this.blobOptions;
+    }
+    return new Blob(this.recordedBlobs, options);
+  }
+
+  animationLoop() {
+    this.ctx.drawArrays(this.ctx.POINTS, 0, 0);
+    if (this.mediaRecorder.state !== "inactive") {
+      requestAnimationFrame(this.animationLoop);
+    }
+  }
+}

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

@@ -6,12 +6,25 @@ import Tooltip from "@mui/material/Tooltip";
 import { useRef, useState } from "react";
 import { vars } from "../../../theme/variables.ts";
 import CustomFormControlLabel from "./CustomFormControlLabel.tsx";
+import { Recorder } from "./Recorder.ts";
+
 const { gray500 } = vars;
 
-function SceneControls({ cameraControlRef, isWireframe, setIsWireframe }) {
+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);
+  const [isRecording, setIsRecording] = useState(false);
+
+  const handleRecordClick = () => {
+    if (isRecording) {
+      stopRecording();
+    } else {
+      startRecording();
+    }
+    setIsRecording(!isRecording);
+  };
+
   const open = Boolean(anchorEl);
   const id = open ? "settings-popover" : undefined;
 
@@ -42,6 +55,24 @@ function SceneControls({ cameraControlRef, isWireframe, setIsWireframe }) {
 
     setIsRotating(!isRotating);
   };
+  const startRecording = () => {
+    if (recorderRef.current === null) {
+      const canvas = document.getElementsByTagName("canvas")[0];
+      recorderRef.current = new Recorder(canvas, {
+        mediaRecorderOptions: { mimeType: "video/webm" },
+        blobOptions: { type: "video/webm" },
+      });
+      recorderRef.current.startRecording();
+    }
+  };
+
+  const stopRecording = async () => {
+    if (recorderRef.current) {
+      recorderRef.current.stopRecording({ type: "video/webm" });
+      recorderRef.current.download("CanvasRecording.webm", { type: "video/webm" });
+      recorderRef.current = null;
+    }
+  };
 
   return (
     <Box
@@ -143,13 +174,17 @@ function SceneControls({ cameraControlRef, isWireframe, setIsWireframe }) {
           <PlayArrowOutlined />
         </IconButton>
       </Tooltip>
-      <Tooltip title="Record viewer" placement="right-start">
-        <IconButton>
-          <RadioButtonCheckedOutlined />
+      <Tooltip title={isRecording ? "Stop recording" : "Record viewer"} placement="right-start">
+        <IconButton onClick={handleRecordClick}>
+          <RadioButtonCheckedOutlined
+            sx={{
+              color: isRecording ? "red" : "inherit",
+            }}
+          />
         </IconButton>
       </Tooltip>
       <Tooltip title="Download graph" placement="right-start">
-        <IconButton>
+        <IconButton onClick={handleScreenshot}>
           <GetAppOutlined />
         </IconButton>
       </Tooltip>

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

@@ -0,0 +1,63 @@
+import * as THREE from "three";
+
+function getResolutionFixedRatio(htmlElement: HTMLElement, target: { width: number; height: number }) {
+  const current = {
+    height: htmlElement.clientHeight,
+    width: htmlElement.clientWidth,
+  };
+
+  if ((Math.abs(target.width - current.width) * 9) / 16 > Math.abs(target.height - current.height)) {
+    return {
+      height: target.height,
+      width: Math.round((current.width * target.height) / current.height),
+    };
+  }
+  return {
+    height: Math.round((current.height * target.width) / current.width),
+    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";
@@ -17,9 +18,10 @@ import {
 import DatasetPicker from "./DatasetPicker.tsx";
 import Gizmo from "./Gizmo.tsx";
 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;
@@ -36,6 +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
@@ -64,10 +71,22 @@ 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"}>
+      <Canvas style={{ backgroundColor: SCENE_BACKGROUND }} frameloop={"demand"} gl={{ preserveDrawingBuffer: false }} onCreated={onCreated}>
+        <color attach="background" args={["#F6F5F4"]} />
         <Suspense fallback={<Loader />}>
           <PerspectiveCamera
             makeDefault
@@ -76,6 +95,7 @@ function ThreeDViewer() {
             position={CAMERA_POSITION}
             near={CAMERA_NEAR}
             far={CAMERA_FAR}
+            ref={cameraRef}
           />
           <CameraControls ref={cameraControlRef} />
 
@@ -87,7 +107,13 @@ function ThreeDViewer() {
           <STLViewer instances={instances} isWireframe={isWireframe} />
         </Suspense>
       </Canvas>
-      <SceneControls cameraControlRef={cameraControlRef} isWireframe={isWireframe} setIsWireframe={setIsWireframe} />
+      <SceneControls
+        cameraControlRef={cameraControlRef}
+        isWireframe={isWireframe}
+        setIsWireframe={setIsWireframe}
+        recorderRef={recorderRef}
+        handleScreenshot={handleScreenshot}
+      />
     </>
   );
 }

applications/visualizer/frontend/src/helpers/utils.ts

@@ -9,3 +9,15 @@ export function areSetsEqual<T>(setA: Set<T>, setB: Set<T>): boolean {
   }
   return true;
 }
+
+export function formatDate(d) {
+  return `${d.getFullYear()}${d.getMonth() + 1}${d.getDate()}-${pad(d.getHours(), 2)}${pad(d.getMinutes(), 2)}${pad(d.getSeconds(), 2)}`;
+}
+
+function pad(num, size) {
+  let s = num + "";
+  while (s.length < size) {
+    s = "0" + s;
+  }
+  return s;
+}