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}