001/* 002 * Copyright (C) Photon Vision. 003 * 004 * This program is free software: you can redistribute it and/or modify 005 * it under the terms of the GNU General Public License as published by 006 * the Free Software Foundation, either version 3 of the License, or 007 * (at your option) any later version. 008 * 009 * This program is distributed in the hope that it will be useful, 010 * but WITHOUT ANY WARRANTY; without even the implied warranty of 011 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 012 * GNU General Public License for more details. 013 * 014 * You should have received a copy of the GNU General Public License 015 * along with this program. If not, see <https://www.gnu.org/licenses/>. 016 */ 017 018package org.photonvision.vision.frame.consumer; 019 020import edu.wpi.first.networktables.IntegerEntry; 021import edu.wpi.first.networktables.IntegerSubscriber; 022import edu.wpi.first.networktables.NetworkTable; 023import edu.wpi.first.networktables.StringSubscriber; 024import edu.wpi.first.wpilibj.DriverStation.MatchType; 025import java.io.File; 026import java.text.DateFormat; 027import java.text.SimpleDateFormat; 028import java.util.Date; 029import java.util.function.Consumer; 030import org.opencv.imgcodecs.Imgcodecs; 031import org.photonvision.common.configuration.ConfigManager; 032import org.photonvision.common.dataflow.networktables.NetworkTablesManager; 033import org.photonvision.common.logging.LogGroup; 034import org.photonvision.common.logging.Logger; 035import org.photonvision.vision.frame.StaticFrames; 036import org.photonvision.vision.opencv.CVMat; 037 038public class FileSaveFrameConsumer implements Consumer<CVMat> { 039 private final Logger logger = new Logger(FileSaveFrameConsumer.class, LogGroup.General); 040 041 // Formatters to generate unique, timestamped file names 042 private static final String FILE_PATH = ConfigManager.getInstance().getImageSavePath().toString(); 043 static final String FILE_EXTENSION = ".jpg"; 044 private static final String NT_SUFFIX = "SaveImgCmd"; 045 046 static final DateFormat df = new SimpleDateFormat("yyyy-MM-dd"); 047 static final DateFormat tf = new SimpleDateFormat("hhmmssSS"); 048 049 private final NetworkTable rootTable; 050 private NetworkTable subTable; 051 private final String ntEntryName; 052 IntegerEntry saveFrameEntry; 053 054 private StringSubscriber ntEventName; 055 private IntegerSubscriber ntMatchNum; 056 private IntegerSubscriber ntMatchType; 057 058 private final String cameraUniqueName; 059 private String cameraNickname; 060 private final String streamType; 061 062 private long savedImagesCount = 0; 063 064 public FileSaveFrameConsumer(String camNickname, String cameraUniqueName, String streamPrefix) { 065 this.ntEntryName = streamPrefix + NT_SUFFIX; 066 this.cameraNickname = camNickname; 067 this.cameraUniqueName = cameraUniqueName; 068 this.streamType = streamPrefix; 069 070 this.rootTable = NetworkTablesManager.getInstance().kRootTable; 071 072 NetworkTable fmsTable = NetworkTablesManager.getInstance().getNTInst().getTable("FMSInfo"); 073 this.ntEventName = fmsTable.getStringTopic("EventName").subscribe("UNKNOWN"); 074 this.ntMatchNum = fmsTable.getIntegerTopic("MatchNumber").subscribe(0); 075 this.ntMatchType = fmsTable.getIntegerTopic("MatchType").subscribe(0); 076 077 updateCameraNickname(camNickname); 078 } 079 080 public void accept(CVMat image) { 081 accept(image, new Date()); 082 } 083 084 public void accept(CVMat image, Date now) { 085 long currentCount = saveFrameEntry.get(); 086 087 // Await save request 088 if (currentCount == -1) return; 089 090 // The requested count is greater than the actual count 091 if (savedImagesCount < currentCount) { 092 String matchData = getMatchData(); 093 094 String fileName = 095 cameraNickname 096 + "_" 097 + streamType 098 + "_" 099 + df.format(now) 100 + "T" 101 + tf.format(now) 102 + "_" 103 + matchData; 104 105 // Check if the Unique Camera directory exists and create it if it doesn't 106 String cameraPath = FILE_PATH + File.separator + this.cameraUniqueName; 107 var cameraDir = new File(cameraPath); 108 if (!cameraDir.exists()) { 109 cameraDir.mkdir(); 110 } 111 112 String saveFilePath = cameraPath + File.separator + fileName + FILE_EXTENSION; 113 114 logger.info("Saving image to: " + saveFilePath); 115 if (image == null || image.getMat() == null || image.getMat().empty()) { 116 Imgcodecs.imwrite(saveFilePath, StaticFrames.LOST_MAT); 117 } else { 118 Imgcodecs.imwrite(saveFilePath, image.getMat()); 119 } 120 121 savedImagesCount++; 122 logger.info("Saved new image at " + saveFilePath); 123 } else if (savedImagesCount > currentCount) { 124 // Reset local value with NT value in case of de-sync 125 savedImagesCount = currentCount; 126 } 127 } 128 129 public void updateCameraNickname(String newCameraNickname) { 130 // Remove existing entries 131 if (this.subTable != null) { 132 if (this.subTable.containsKey(ntEntryName)) { 133 this.subTable.getEntry(ntEntryName).close(); 134 } 135 } 136 137 // Recreate and re-init network tables structure 138 this.cameraNickname = newCameraNickname; 139 this.subTable = rootTable.getSubTable(this.cameraNickname); 140 this.subTable.getEntry(ntEntryName).setInteger(savedImagesCount); 141 this.saveFrameEntry = subTable.getIntegerTopic(ntEntryName).getEntry(-1); // Default negative 142 } 143 144 public void overrideTakeSnapshot() { 145 // Simulate NT change 146 saveFrameEntry.set(saveFrameEntry.get() + 1); 147 } 148 149 /** 150 * Returns the match Data collected from the NT. eg : Q58 for qualification match 58. If not in 151 * event, returns None-0-Unknown 152 */ 153 private String getMatchData() { 154 var matchType = ntMatchType.getAtomic(); 155 if (matchType.timestamp == 0) { 156 // no NT info yet 157 logger.warn("Did not receive match type, defaulting to None"); 158 } 159 160 var matchNum = ntMatchNum.getAtomic(); 161 if (matchNum.timestamp == 0) { 162 logger.warn("Did not receive match num, defaulting to 0"); 163 } 164 165 var eventName = ntEventName.getAtomic(); 166 if (eventName.timestamp == 0) { 167 logger.warn("Did not receive event name, defaulting to 'UNKNOWN'"); 168 } 169 170 MatchType wpiMatchType = MatchType.None; // Default is to be unknown 171 if (matchType.value < 0 || matchType.value >= MatchType.values().length) { 172 logger.error("Invalid match type from FMS: " + matchType.value); 173 } else { 174 wpiMatchType = MatchType.values()[(int) matchType.value]; 175 } 176 177 return wpiMatchType.name() + "-" + matchNum.value + "-" + eventName.value; 178 } 179 180 public void close() { 181 // troododfa;lkjadsf;lkfdsaj otgooadflsk;j 182 } 183}