Refactor BakedHeatMap and useHeatMapStore to support typed baked points; enhance heatmap export functionality

app/src/components/heatMapGenerator/baked/bakedHeatMap.tsx

@@ -13,33 +13,39 @@ const BakedHeatMap = () => {
     const { bakedPoints } = useHeatMapStore();
     const materialRef = useRef<THREE.ShaderMaterial>(null);
     const meshRef = useRef<THREE.Mesh>(null);
-    const { gl, scene, camera, size } = useThree();
+    const { gl } = useThree();
 
     const height = CONSTANTS.gridConfig.size;
     const width = CONSTANTS.gridConfig.size;
 
-    const pointTexture = useMemo(() => {
-        if (bakedPoints.length === 0) return null;
+    /**
+     * Helper: Create a DataTexture from filtered baked points
+     */
+    const createPointTexture = useCallback(
+        (filteredPoints: typeof bakedPoints) => {
+            if (filteredPoints.length === 0) return null;
 
-        const data = new Float32Array(bakedPoints.length * 4);
-        bakedPoints.forEach((p, i) => {
-            const index = i * 4;
-            data[index] = (p.x + width / 2) / width;
-            data[index + 1] = (p.y + height / 2) / height;
-            data[index + 2] = 0.3;
-            data[index + 3] = 0.0;
-        });
+            const data = new Float32Array(filteredPoints.length * 4);
+            filteredPoints.forEach((p, i) => {
+                const index = i * 4;
+                data[index] = (p.points.x + width / 2) / width;
+                data[index + 1] = (p.points.y + height / 2) / height;
+                data[index + 2] = 0.3; // heat strength
+                data[index + 3] = 0.0; // unused
+            });
 
-        const texture = new THREE.DataTexture(
-            data,
-            bakedPoints.length,
-            1,
-            THREE.RGBAFormat,
-            THREE.FloatType
-        );
-        texture.needsUpdate = true;
-        return texture;
-    }, [bakedPoints, width, height]);
+            const texture = new THREE.DataTexture(
+                data,
+                filteredPoints.length,
+                1,
+                THREE.RGBAFormat,
+                THREE.FloatType
+            );
+            texture.needsUpdate = true;
+            return texture;
+        },
+        [width, height]
+    );
 
     const uniformsRef = useRef({
         u_points: { value: null as THREE.DataTexture | null },
@@ -49,66 +55,98 @@ const BakedHeatMap = () => {
         u_growthRate: { value: GROWTH_RATE },
     });
 
+    /**
+     * Render the heatmap to a base64 PNG for a specific type (human/vehicle)
+     */
+    const renderHeatmapToImage = useCallback(
+        (type: string) => {
+            if (!meshRef.current) return null;
+
+            // Filter points by type first
+            const filteredPoints = bakedPoints.filter((p) => p.type === type);
+            if (filteredPoints.length === 0) return null;
+
+            // Create texture for these points
+            const pointTexture = createPointTexture(filteredPoints);
+            if (!pointTexture) return null;
+
+            uniformsRef.current.u_points.value = pointTexture;
+            uniformsRef.current.u_count.value = filteredPoints.length;
+
+            // Temporary top-down orthographic camera
+            const exportCamera = new THREE.OrthographicCamera(
+                width / -2,
+                width / 2,
+                height / 2,
+                height / -2,
+                0.1,
+                10
+            );
+            exportCamera.position.set(0, 1, 0);
+            exportCamera.lookAt(0, 0, 0);
+
+            // Offscreen render target
+            const renderTarget = new THREE.WebGLRenderTarget(1024, 1024, {
+                format: THREE.RGBAFormat,
+                type: THREE.UnsignedByteType,
+            });
+
+            // Temporary scene with only this mesh
+            const tempScene = new THREE.Scene();
+            tempScene.add(meshRef.current);
+
+            // Render to texture
+            gl.setRenderTarget(renderTarget);
+            gl.render(tempScene, exportCamera);
+            gl.setRenderTarget(null);
+
+            // Read pixels
+            const pixels = new Uint8Array(1024 * 1024 * 4);
+            gl.readRenderTargetPixels(renderTarget, 0, 0, 1024, 1024, pixels);
+
+            // Convert pixels to base64 PNG
+            const canvas = document.createElement("canvas");
+            canvas.width = 1024;
+            canvas.height = 1024;
+            const ctx = canvas.getContext("2d");
+
+            if (ctx) {
+                const imageData = ctx.createImageData(1024, 1024);
+                imageData.data.set(pixels);
+                ctx.putImageData(imageData, 0, 0);
+                return canvas.toDataURL("image/png");
+            }
+            return null;
+        },
+        [gl, width, height, bakedPoints, createPointTexture]
+    );
+
+    /**
+     * Export both human and vehicle heatmaps and log them
+     */
+    const exportHeatmapAsPNG = useCallback(() => {
+        const types = ["human", "vehicle"];
+
+        const result = types.map((type) => {
+            const image = renderHeatmapToImage(type);
+            return {
+                type,
+                image,
+            };
+        });
+
+        console.log("Exported Heatmaps:", result);
+        return result;
+    }, [renderHeatmapToImage]);
+
+    // Create default texture for initial rendering
+    const pointTexture = useMemo(() => createPointTexture(bakedPoints), [bakedPoints, createPointTexture]);
+
     useEffect(() => {
         uniformsRef.current.u_points.value = pointTexture;
         uniformsRef.current.u_count.value = bakedPoints.length;
     }, [pointTexture, bakedPoints.length]);
 
-    /**
-     * Export the heatmap as PNG
-     */
-    const exportHeatmapAsPNG = useCallback(() => {
-        if (!meshRef.current) return;
-
-        // Create a temporary orthographic camera
-        const exportCamera = new THREE.OrthographicCamera(
-            width / -2,
-            width / 2,
-            height / 2,
-            height / -2,
-            0.1,
-            10
-        );
-        exportCamera.position.set(0, 1, 0); // top-down view
-        exportCamera.lookAt(0, 0, 0);
-
-        // Create an offscreen render target
-        const renderTarget = new THREE.WebGLRenderTarget(1024, 1024, {
-            format: THREE.RGBAFormat,
-            type: THREE.UnsignedByteType,
-        });
-
-        // Render only the heatmap mesh
-        const tempScene = new THREE.Scene();
-        tempScene.add(meshRef.current);
-
-        gl.setRenderTarget(renderTarget);
-        gl.render(tempScene, exportCamera);
-        gl.setRenderTarget(null);
-
-        // Read pixels
-        const pixels = new Uint8Array(1024 * 1024 * 4);
-        gl.readRenderTargetPixels(renderTarget, 0, 0, 1024, 1024, pixels);
-
-        // Convert pixel data to image
-        const canvas = document.createElement("canvas");
-        canvas.width = 1024;
-        canvas.height = 1024;
-        const ctx = canvas.getContext("2d");
-
-        if (ctx) {
-            const imageData = ctx.createImageData(1024, 1024);
-            imageData.data.set(pixels);
-            ctx.putImageData(imageData, 0, 0);
-
-            // Create a download link
-            const link = document.createElement("a");
-            link.href = canvas.toDataURL("image/png");
-            link.download = "heatmap.png";
-            link.click();
-        }
-    }, [gl, width, height]);
-
     return (
         <>
             <mesh ref={meshRef} rotation={[Math.PI / 2, 0, 0]} position={[0, 0.025, 0]}>
@@ -173,24 +211,24 @@ const BakedHeatMap = () => {
                 />
             </mesh>
 
-            {/* Button to trigger export */}
+            {/* Button to export */}
             <Html>
-            <button
-                style={{
-                    position: "absolute",
-                    top: "10px",
-                    left: "10px",
-                    padding: "8px 12px",
-                    background: "#333",
-                    color: "white",
-                    border: "none",
-                    cursor: "pointer",
-                    zIndex: 10,
-                }}
-                onClick={exportHeatmapAsPNG}
-            >
-                Export Heatmap
-            </button>
+                <button
+                    style={{
+                        position: "absolute",
+                        top: "10px",
+                        left: "10px",
+                        padding: "8px 12px",
+                        background: "#333",
+                        color: "white",
+                        border: "none",
+                        cursor: "pointer",
+                        zIndex: 10,
+                    }}
+                    onClick={exportHeatmapAsPNG}
+                >
+                    Export Heatmap
+                </button>
             </Html>
         </>
     );

app/src/components/heatMapGenerator/heatMap.tsx

@@ -7,7 +7,7 @@ function HeatMap() {
     const { bakedPoints, setBakedPoints } = useHeatMapStore();
 
     useEffect(() => {
-        // console.log("bakedPoints: ", bakedPoints);
+        console.log("bakedPoints: ", bakedPoints);
     }, [bakedPoints]);
 
     return (

app/src/components/heatMapGenerator/realTime/realTimeHeatMap.tsx

@@ -126,7 +126,10 @@ const RealTimeHeatMap = () => {
                   })()
                 : model.position;
 
-            addBakedPoint({ x: pos.x, y: pos.z });
+            addBakedPoint({
+                type: event.type,
+                points: { x: pos.x, y: pos.z },
+            });
 
             updatedPoints.push({ x: pos.x, y: pos.z, strength: 0.3, lastUpdated: now });
         });

app/src/store/simulation/useHeatMapStore.ts

@@ -6,27 +6,37 @@ interface BakedPoint {
     y: number;
 }
 
+interface BakedPointImage {
+    type: string;
+    points: BakedPoint;
+}
+
 interface HeatMapState {
     monitoringVehicle: string[];
     monitoringHuman: string[];
-    bakedPoints: BakedPoint[];
+    bakedPoints: BakedPointImage[];
 
+    // Vehicle monitoring
     setMonitoringVehicle: (vehiclesUuid: string[]) => void;
     addMonitoringVehicle: (vehicleUuid: string) => void;
     removeMonitoringVehicle: (vehicleUuid: string) => void;
     clearMonitoringVehicle: () => void;
 
+    // Human monitoring
     setMonitoringHuman: (humansUuid: string[]) => void;
     addMonitoringHuman: (humanUuid: string) => void;
     removeMonitoringHuman: (humanUuid: string) => void;
     clearMonitoringHuman: () => void;
 
+    // Clear both monitors
     clearAllMonitors: () => void;
 
-    setBakedPoints: (points: BakedPoint[]) => void;
-    addBakedPoint: (point: BakedPoint) => void;
+    // Baked points
+    setBakedPoints: (points: BakedPointImage[]) => void;
+    addBakedPoint: (point: BakedPointImage) => void;
     clearBakedPoints: () => void;
 
+    // Utility checkers
     hasVehicle: (vehicleUuid: string) => boolean;
     hasHuman: (humanUuid: string) => boolean;
 }
@@ -37,6 +47,7 @@ export const useHeatMapStore = create<HeatMapState>()(
         monitoringHuman: [],
         bakedPoints: [],
 
+        // Vehicles
         setMonitoringVehicle: (vehiclesUuid) =>
             set((state) => {
                 state.monitoringVehicle = vehiclesUuid;
@@ -59,6 +70,7 @@ export const useHeatMapStore = create<HeatMapState>()(
                 state.monitoringVehicle = [];
             }),
 
+        // Humans
         setMonitoringHuman: (humansUuid) =>
             set((state) => {
                 state.monitoringHuman = humansUuid;
@@ -81,17 +93,20 @@ export const useHeatMapStore = create<HeatMapState>()(
                 state.monitoringHuman = [];
             }),
 
+        // Clear all
         clearAllMonitors: () =>
             set((state) => {
                 state.monitoringVehicle = [];
                 state.monitoringHuman = [];
             }),
 
-        setBakedPoints: (points: BakedPoint[]) =>
+        // Baked Points
+        setBakedPoints: (points) =>
             set((state) => {
                 state.bakedPoints = points;
             }),
-        addBakedPoint: (point: BakedPoint) =>
+
+        addBakedPoint: (point) =>
             set((state) => {
                 state.bakedPoints.push(point);
             }),
@@ -101,6 +116,7 @@ export const useHeatMapStore = create<HeatMapState>()(
                 state.bakedPoints = [];
             }),
 
+        // Utility checkers
         hasVehicle: (vehicleUuid) => get().monitoringVehicle.includes(vehicleUuid),
         hasHuman: (humanUuid) => get().monitoringHuman.includes(humanUuid),
     }))