From ededc4f1302b2be251f15c1d15604e4e056e8d41 Mon Sep 17 00:00:00 2001
From: Sriman Achanta <68172138+srimanachanta@users.noreply.github.com>
Date: Tue, 17 Oct 2023 09:08:25 -0400
Subject: [PATCH] Take Snapshots from PhotonClient (#940)

---
 .../components/app/photon-camera-stream.vue   | 74 ++++++++++++++---
 .../cameras/CameraCalibrationCard.vue         |  2 +-
 .../src/components/cameras/CamerasView.vue    |  6 +-
 .../stores/settings/CameraSettingsStore.ts    | 29 ++++++-
 .../frame/consumer/FileSaveFrameConsumer.java | 80 ++++++++++---------
 .../vision/processes/VisionModule.java        |  8 ++
 .../VisionModuleChangeSubscriber.java         |  6 ++
 .../server/DataSocketHandler.java             | 24 ++++++
 .../server/DataSocketMessageType.java         |  2 +
 9 files changed, 174 insertions(+), 57 deletions(-)

diff --git a/photon-client/src/components/app/photon-camera-stream.vue b/photon-client/src/components/app/photon-camera-stream.vue
index 2868ae718..1349da2f9 100644
--- a/photon-client/src/components/app/photon-camera-stream.vue
+++ b/photon-client/src/components/app/photon-camera-stream.vue
@@ -4,13 +4,14 @@ import { useCameraSettingsStore } from "@/stores/settings/CameraSettingsStore";
 import { useStateStore } from "@/stores/StateStore";
 import loadingImage from "@/assets/images/loading.svg";
 import type { StyleValue } from "vue/types/jsx";
+import CvIcon from "@/components/common/cv-icon.vue";
 
 const props = defineProps<{
   streamType: "Raw" | "Processed";
   id?: string;
 }>();
 
-const src = computed<string>(() => {
+const streamSrc = computed<string>(() => {
   const port =
     useCameraSettingsStore().currentCameraSettings.stream[props.streamType === "Raw" ? "inputPort" : "outputPort"];
 
@@ -20,25 +21,74 @@ const src = computed<string>(() => {
 
   return `http://${inject("backendHostname")}:${port}/stream.mjpg`;
 });
-const alt = computed<string>(() => `${props.streamType} Stream View`);
-
-const style = computed<StyleValue>(() => {
+const streamDesc = computed<string>(() => `${props.streamType} Stream View`);
+const streamStyle = computed<StyleValue>(() => {
   if (useStateStore().colorPickingMode) {
-    return { cursor: "crosshair" };
-  } else if (src.value !== loadingImage) {
-    return { cursor: "pointer" };
+    return { width: "100%", cursor: "crosshair" };
+  } else if (streamSrc.value !== loadingImage) {
+    return { width: "100%", cursor: "pointer" };
   }
 
-  return {};
+  return { width: "100%" };
 });
 
-const handleClick = () => {
-  if (!useStateStore().colorPickingMode && src.value !== loadingImage) {
-    window.open(src.value);
+const overlayStyle = computed<StyleValue>(() => {
+  if (useStateStore().colorPickingMode || streamSrc.value == loadingImage) {
+    return { display: "none" };
+  } else {
+    return {};
+  }
+});
+
+const handleStreamClick = () => {
+  if (!useStateStore().colorPickingMode && streamSrc.value !== loadingImage) {
+    window.open(streamSrc.value);
+  }
+};
+const handleCaptureClick = () => {
+  if (props.streamType === "Raw") {
+    useCameraSettingsStore().saveInputSnapshot();
+  } else {
+    useCameraSettingsStore().saveOutputSnapshot();
   }
 };
 </script>
 
 <template>
-  <img :id="id" crossorigin="anonymous" :src="src" :alt="alt" :style="style" @click="handleClick" />
+  <div class="stream-container">
+    <img
+      :id="id"
+      crossorigin="anonymous"
+      :src="streamSrc"
+      :alt="streamDesc"
+      :style="streamStyle"
+      @click="handleStreamClick"
+    />
+    <div class="stream-overlay" :style="overlayStyle">
+      <cv-icon
+        icon-name="mdi-camera-image"
+        tooltip="Capture and save a frame of this stream"
+        class="ma-1 mr-2"
+        @click="handleCaptureClick"
+      />
+    </div>
+  </div>
 </template>
+
+<style scoped>
+.stream-container {
+  position: relative;
+}
+
+.stream-overlay {
+  opacity: 0;
+  transition: 0.1s ease;
+  position: absolute;
+  top: 0;
+  right: 0;
+}
+
+.stream-container:hover .stream-overlay {
+  opacity: 1;
+}
+</style>
diff --git a/photon-client/src/components/cameras/CameraCalibrationCard.vue b/photon-client/src/components/cameras/CameraCalibrationCard.vue
index fb1222465..969dffa86 100644
--- a/photon-client/src/components/cameras/CameraCalibrationCard.vue
+++ b/photon-client/src/components/cameras/CameraCalibrationCard.vue
@@ -384,7 +384,7 @@ const endCalibration = () => {
               color="secondary"
               style="width: 100%"
               :disabled="!settingsValid"
-              @click="isCalibrating ? useCameraSettingsStore().takeCalibrationSnapshot(true) : startCalibration()"
+              @click="isCalibrating ? useCameraSettingsStore().takeCalibrationSnapshot() : startCalibration()"
             >
               {{ isCalibrating ? "Take Snapshot" : "Start Calibration" }}
             </v-btn>
diff --git a/photon-client/src/components/cameras/CamerasView.vue b/photon-client/src/components/cameras/CamerasView.vue
index 8b7f1785d..c9e3d32cb 100644
--- a/photon-client/src/components/cameras/CamerasView.vue
+++ b/photon-client/src/components/cameras/CamerasView.vue
@@ -41,7 +41,7 @@ const fpsTooLow = computed<boolean>(() => {
 </script>
 
 <template>
-  <v-card class="mb-3 pr-6 pb-3 pa-4" color="primary" dark>
+  <v-card class="mb-3 pb-3 pa-4" color="primary" dark>
     <v-card-title
       class="pb-0 mb-2 pl-4 pt-1"
       style="min-height: 50px; justify-content: space-between; align-content: center"
@@ -133,16 +133,16 @@ th {
 .stream {
   display: flex;
   justify-content: center;
-  width: 100%;
 }
 
 @media only screen and (min-width: 512px) and (max-width: 960px) {
   .stream-container {
     flex-wrap: nowrap;
+    justify-content: center;
   }
 
   .stream {
-    width: 50%;
+    max-width: 50%;
   }
 }
 </style>
diff --git a/photon-client/src/stores/settings/CameraSettingsStore.ts b/photon-client/src/stores/settings/CameraSettingsStore.ts
index 673862771..356d12faa 100644
--- a/photon-client/src/stores/settings/CameraSettingsStore.ts
+++ b/photon-client/src/stores/settings/CameraSettingsStore.ts
@@ -357,12 +357,35 @@ export const useCameraSettingsStore = defineStore("cameraSettings", {
     /**
      * Take a snapshot for the calibration processes
      *
-     * @param takeSnapshot whether or not to take a snapshot. Defaults to true
      * @param cameraIndex the index of the camera that is currently in the calibration process
      */
-    takeCalibrationSnapshot(takeSnapshot = true, cameraIndex: number = useStateStore().currentCameraIndex) {
+    takeCalibrationSnapshot(cameraIndex: number = useStateStore().currentCameraIndex) {
       const payload = {
-        takeCalibrationSnapshot: takeSnapshot,
+        takeCalibrationSnapshot: true,
+        cameraIndex: cameraIndex
+      };
+      useStateStore().websocket?.send(payload, true);
+    },
+    /**
+     * Save a snapshot of the input frame of the camera.
+     *
+     * @param cameraIndex the index of the camera
+     */
+    saveInputSnapshot(cameraIndex: number = useStateStore().currentCameraIndex) {
+      const payload = {
+        saveInputSnapshot: true,
+        cameraIndex: cameraIndex
+      };
+      useStateStore().websocket?.send(payload, true);
+    },
+    /**
+     * Save a snapshot of the output frame of the camera.
+     *
+     * @param cameraIndex the index of the camera
+     */
+    saveOutputSnapshot(cameraIndex: number = useStateStore().currentCameraIndex) {
+      const payload = {
+        saveOutputSnapshot: true,
         cameraIndex: cameraIndex
       };
       useStateStore().websocket?.send(payload, true);
diff --git a/photon-core/src/main/java/org/photonvision/vision/frame/consumer/FileSaveFrameConsumer.java b/photon-core/src/main/java/org/photonvision/vision/frame/consumer/FileSaveFrameConsumer.java
index 2552b5fbd..57bc312b2 100644
--- a/photon-core/src/main/java/org/photonvision/vision/frame/consumer/FileSaveFrameConsumer.java
+++ b/photon-core/src/main/java/org/photonvision/vision/frame/consumer/FileSaveFrameConsumer.java
@@ -32,58 +32,57 @@ import org.photonvision.common.logging.Logger;
 import org.photonvision.vision.opencv.CVMat;
 
 public class FileSaveFrameConsumer implements Consumer<CVMat> {
+    private final Logger logger = new Logger(FileSaveFrameConsumer.class, LogGroup.General);
+
     // Formatters to generate unique, timestamped file names
     private static final String FILE_PATH = ConfigManager.getInstance().getImageSavePath().toString();
     private static final String FILE_EXTENSION = ".jpg";
+    private static final String NT_SUFFIX = "SaveImgCmd";
+
     DateFormat df = new SimpleDateFormat("yyyy-MM-dd");
     DateFormat tf = new SimpleDateFormat("hhmmssSS");
-    private final String NT_SUFFIX = "SaveImgCmd";
-    private final String ntEntryName;
-    private NetworkTable subTable;
+
     private final NetworkTable rootTable;
-    private final Logger logger;
-    private long imgSaveCountInternal = 0;
-    private String camNickname;
-    private final String fnamePrefix;
-    private IntegerEntry entry;
+    private NetworkTable subTable;
+    private final String ntEntryName;
+    private IntegerEntry saveFrameEntry;
+
+    private String cameraNickname;
+    private final String streamType;
+
+    private long savedImagesCount = 0;
 
     public FileSaveFrameConsumer(String camNickname, String streamPrefix) {
-        this.fnamePrefix = camNickname + "_" + streamPrefix;
         this.ntEntryName = streamPrefix + NT_SUFFIX;
+        this.cameraNickname = camNickname;
+        this.streamType = streamPrefix;
+
         this.rootTable = NetworkTablesManager.getInstance().kRootTable;
         updateCameraNickname(camNickname);
-        this.logger = new Logger(FileSaveFrameConsumer.class, this.camNickname, LogGroup.General);
     }
 
     public void accept(CVMat image) {
         if (image != null && image.getMat() != null && !image.getMat().empty()) {
-            var curCommand = entry.get(); // default to just our current count
-            if (curCommand >= 0) {
-                // Only do something if we got a valid current command
-                if (imgSaveCountInternal < curCommand) {
-                    // Save one frame.
-                    // Create the filename
-                    Date now = new Date();
-                    String savefile =
-                            FILE_PATH
-                                    + File.separator
-                                    + fnamePrefix
-                                    + "_"
-                                    + df.format(now)
-                                    + "T"
-                                    + tf.format(now)
-                                    + FILE_EXTENSION;
+            long currentCount = saveFrameEntry.get();
 
-                    // write to file
-                    Imgcodecs.imwrite(savefile, image.getMat());
+            // Await save request
+            if (currentCount == -1) return;
 
-                    // Count one more image saved
-                    imgSaveCountInternal++;
-                    logger.info("Saved new image at " + savefile);
+            // The requested count is greater than the actual count
+            if (savedImagesCount < currentCount) {
+                Date now = new Date();
 
-                } else if (imgSaveCountInternal > curCommand) {
-                    imgSaveCountInternal = curCommand;
-                }
+                String fileName =
+                        cameraNickname + "_" + streamType + "_" + df.format(now) + "T" + tf.format(now);
+                String saveFilePath = FILE_PATH + File.separator + fileName + FILE_EXTENSION;
+
+                Imgcodecs.imwrite(saveFilePath, image.getMat());
+
+                savedImagesCount++;
+                logger.info("Saved new image at " + saveFilePath);
+            } else if (savedImagesCount > currentCount) {
+                // Reset local value with NT value in case of de-sync
+                savedImagesCount = currentCount;
             }
         }
     }
@@ -97,9 +96,14 @@ public class FileSaveFrameConsumer implements Consumer<CVMat> {
         }
 
         // Recreate and re-init network tables structure
-        this.camNickname = newCameraNickname;
-        this.subTable = rootTable.getSubTable(this.camNickname);
-        this.subTable.getEntry(ntEntryName).setInteger(imgSaveCountInternal);
-        this.entry = subTable.getIntegerTopic(ntEntryName).getEntry(-1); // Default negative
+        this.cameraNickname = newCameraNickname;
+        this.subTable = rootTable.getSubTable(this.cameraNickname);
+        this.subTable.getEntry(ntEntryName).setInteger(savedImagesCount);
+        this.saveFrameEntry = subTable.getIntegerTopic(ntEntryName).getEntry(-1); // Default negative
+    }
+
+    public void overrideTakeSnapshot() {
+        // Simulate NT change
+        saveFrameEntry.set(saveFrameEntry.get() + 1);
     }
 }
diff --git a/photon-core/src/main/java/org/photonvision/vision/processes/VisionModule.java b/photon-core/src/main/java/org/photonvision/vision/processes/VisionModule.java
index d246a45b5..75978739e 100644
--- a/photon-core/src/main/java/org/photonvision/vision/processes/VisionModule.java
+++ b/photon-core/src/main/java/org/photonvision/vision/processes/VisionModule.java
@@ -354,6 +354,14 @@ public class VisionModule {
         setPipeline(PipelineManager.CAL_3D_INDEX);
     }
 
+    public void saveInputSnapshot() {
+        inputFrameSaver.overrideTakeSnapshot();
+    }
+
+    public void saveOutputSnapshot() {
+        outputFrameSaver.overrideTakeSnapshot();
+    }
+
     public void takeCalibrationSnapshot() {
         pipelineManager.calibration3dPipeline.takeSnapshot();
     }
diff --git a/photon-core/src/main/java/org/photonvision/vision/processes/VisionModuleChangeSubscriber.java b/photon-core/src/main/java/org/photonvision/vision/processes/VisionModuleChangeSubscriber.java
index e16e4eeed..e824ce89d 100644
--- a/photon-core/src/main/java/org/photonvision/vision/processes/VisionModuleChangeSubscriber.java
+++ b/photon-core/src/main/java/org/photonvision/vision/processes/VisionModuleChangeSubscriber.java
@@ -107,6 +107,12 @@ public class VisionModuleChangeSubscriber extends DataChangeSubscriber {
                         parentModule.startCalibration(data);
                         parentModule.saveAndBroadcastAll();
                         return;
+                    case "saveInputSnapshot":
+                        parentModule.saveInputSnapshot();
+                        return;
+                    case "saveOutputSnapshot":
+                        parentModule.saveOutputSnapshot();
+                        return;
                     case "takeCalSnapshot":
                         parentModule.takeCalibrationSnapshot();
                         return;
diff --git a/photon-server/src/main/java/org/photonvision/server/DataSocketHandler.java b/photon-server/src/main/java/org/photonvision/server/DataSocketHandler.java
index 19fc18f93..f7f034572 100644
--- a/photon-server/src/main/java/org/photonvision/server/DataSocketHandler.java
+++ b/photon-server/src/main/java/org/photonvision/server/DataSocketHandler.java
@@ -265,6 +265,30 @@ public class DataSocketHandler {
                                 dcService.publishEvent(changePipelineEvent);
                                 break;
                             }
+                        case SMT_SAVEINPUTSNAPSHOT:
+                            {
+                                var takeInputSnapshotEvent =
+                                        new IncomingWebSocketEvent<>(
+                                                DataChangeDestination.DCD_ACTIVEMODULE,
+                                                "saveInputSnapshot",
+                                                0,
+                                                cameraIndex,
+                                                context);
+                                dcService.publishEvent(takeInputSnapshotEvent);
+                                break;
+                            }
+                        case SMT_SAVEOUTPUTSNAPSHOT:
+                            {
+                                var takeOutputSnapshotEvent =
+                                        new IncomingWebSocketEvent<>(
+                                                DataChangeDestination.DCD_ACTIVEMODULE,
+                                                "saveOutputSnapshot",
+                                                0,
+                                                cameraIndex,
+                                                context);
+                                dcService.publishEvent(takeOutputSnapshotEvent);
+                                break;
+                            }
                         case SMT_TAKECALIBRATIONSNAPSHOT:
                             {
                                 var takeCalSnapshotEvent =
diff --git a/photon-server/src/main/java/org/photonvision/server/DataSocketMessageType.java b/photon-server/src/main/java/org/photonvision/server/DataSocketMessageType.java
index 256695f2c..f5222cf37 100644
--- a/photon-server/src/main/java/org/photonvision/server/DataSocketMessageType.java
+++ b/photon-server/src/main/java/org/photonvision/server/DataSocketMessageType.java
@@ -32,6 +32,8 @@ public enum DataSocketMessageType {
     SMT_PIPELINESETTINGCHANGE("changePipelineSetting"),
     SMT_CURRENTPIPELINE("currentPipeline"),
     SMT_STARTPNPCALIBRATION("startPnpCalibration"),
+    SMT_SAVEINPUTSNAPSHOT("saveInputSnapshot"),
+    SMT_SAVEOUTPUTSNAPSHOT("saveOutputSnapshot"),
     SMT_TAKECALIBRATIONSNAPSHOT("takeCalibrationSnapshot"),
     SMT_DUPLICATEPIPELINE("duplicatePipeline"),
     SMT_CHANGEBRIGHTNESS("enabledLEDPercentage"),