PhotonVision/Main/src/main/java/com/chameleonvision/vision/process/CameraProcess.java

package com.chameleonvision.vision.process;

import com.chameleonvision.MemoryManager;
import com.chameleonvision.settings.SettingsManager;
import com.chameleonvision.vision.Pipeline;
import com.chameleonvision.vision.camera.Camera;
import com.chameleonvision.vision.camera.CameraValues;
import com.chameleonvision.web.Server;
import edu.wpi.cscore.CvSink;
import edu.wpi.cscore.CvSource;
import edu.wpi.first.cameraserver.CameraServer;
import edu.wpi.first.networktables.*;
import org.opencv.core.*;
import org.opencv.imgproc.Imgproc;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;

public class CameraProcess implements Runnable {

    private final Camera camera;
    private final String cameraName;

    // NetworkTables
    private NetworkTableEntry ntPipelineEntry;
    private NetworkTableEntry ntDriverModeEntry;
    private NetworkTableEntry ntYawEntry;
    private NetworkTableEntry ntPitchEntry;
    private NetworkTableEntry ntDistanceEntry;
    private NetworkTableEntry ntTimeStampEntry;
    private NetworkTableEntry ntValidEntry;

    private MemoryManager memManager = new MemoryManager(125);

    // chameleon specific
    private Pipeline currentPipeline;
    private VisionProcess visionProcess;
    private CameraValues camVals;

    // pipeline process items
    private List<MatOfPoint> FoundContours = new ArrayList<>();
    private List<MatOfPoint> FilteredContours = new ArrayList<>();
    private List<RotatedRect> GroupedContours = new ArrayList<>();
    private Mat cameraInputMat = new Mat();
    private Mat hsvThreshMat = new Mat();
    private Mat streamOutputMat = new Mat();
    private Scalar contourRectColor = new Scalar(255, 0, 0);
    private long TimeStamp = 0;

    private void DriverModeListener(EntryNotification entryNotification) {
        if (entryNotification.value.getBoolean()) {
            camera.setExposure(25);
            camera.setBrightness(15);
        } else {
            Pipeline pipeline = camera.getCurrentPipeline();
            camera.setExposure(pipeline.exposure);
            camera.setBrightness(pipeline.brightness);
        }
    }

    private void PipelineListener(EntryNotification entryNotification) {
        if (camera.getPipelines().containsKey(entryNotification.value.getString())) {
//            camera.setPntryNotification.value.getString());
            Pipeline pipeline = camera.getCurrentPipeline();

            camera.setExposure(pipeline.exposure);
            camera.setBrightness(pipeline.brightness);
            //TODO Send Pipeline change using websocket to client
        } else {
            ntPipelineEntry.setString("pipeline" + camera.getCurrentPipelineIndex());
        }
    }

    public CameraProcess(Camera processCam) {
        camera = processCam;
        this.cameraName = camera.name;

        // NetworkTables
        NetworkTable ntTable = NetworkTableInstance.getDefault().getTable("/chameleon-vision/" + cameraName);
        ntPipelineEntry = ntTable.getEntry("Pipeline");
        ntDriverModeEntry = ntTable.getEntry("Driver_Mode");
        ntPitchEntry = ntTable.getEntry("Pitch");
        ntYawEntry = ntTable.getEntry("Yaw");
        ntDistanceEntry = ntTable.getEntry("Distance");
        ntTimeStampEntry = ntTable.getEntry("TimeStamp");
        ntValidEntry = ntTable.getEntry("Valid");
        ntDriverModeEntry.addListener(this::DriverModeListener, EntryListenerFlags.kUpdate);
        ntPipelineEntry.addListener(this::PipelineListener, EntryListenerFlags.kUpdate);
        ntDriverModeEntry.setBoolean(false);
        ntPipelineEntry.setString("pipeline" + camera.getCurrentPipelineIndex());

        // camera settings
        camVals = new CameraValues(camera);
        visionProcess = new VisionProcess(camVals);
    }

    private void drawContour(Mat inputMat, RotatedRect contourRect) {
        if (contourRect == null) return;
        List<MatOfPoint> drawnContour = new ArrayList<>();
        Point[] vertices = new Point[4];
        contourRect.points(vertices);
        drawnContour.add(new MatOfPoint(vertices));
        Imgproc.drawContours(inputMat, drawnContour, 0, contourRectColor, 3);
        Imgproc.circle(inputMat, contourRect.center, 3, contourRectColor);
    }

    private void updateNetworkTables(PipelineResult pipelineResult) {
        ntValidEntry.setBoolean(pipelineResult.IsValid);
        if (pipelineResult.IsValid){
            ntYawEntry.setNumber(pipelineResult.Yaw);
            ntPitchEntry.setNumber(pipelineResult.Pitch);
        }
        ntTimeStampEntry.setNumber(TimeStamp);
    }

    // TODO: Separate video output to separate function, maybe even second thread
    private PipelineResult runVisionProcess(Mat inputImage, Mat outputImage) {
        var pipelineResult = new PipelineResult();

        if (currentPipeline == null) {
            return pipelineResult;
        }
        if (ntDriverModeEntry.getBoolean(false)){
            inputImage.copyTo(outputImage);
            return pipelineResult;
        }
        Scalar hsvLower = new Scalar(currentPipeline.hue.get(0), currentPipeline.saturation.get(0), currentPipeline.value.get(0));
        Scalar hsvUpper = new Scalar(currentPipeline.hue.get(1), currentPipeline.saturation.get(1), currentPipeline.value.get(1));

        visionProcess.HSVThreshold(inputImage, hsvThreshMat, hsvLower, hsvUpper, currentPipeline.erode, currentPipeline.dilate);

        if (currentPipeline.is_binary == 1) {
            Imgproc.cvtColor(hsvThreshMat, outputImage, Imgproc.COLOR_GRAY2BGR, 3);
        } else {
            inputImage.copyTo(outputImage);
        }
        FoundContours = visionProcess.FindContours(hsvThreshMat);
        if (FoundContours.size() > 0) {
            FilteredContours = visionProcess.FilterContours(FoundContours, currentPipeline.area, currentPipeline.ratio, currentPipeline.extent, currentPipeline.sort_mode, currentPipeline.target_intersection, currentPipeline.target_group);
            if (FilteredContours.size() > 0) {
                GroupedContours = visionProcess.GroupTargets(FilteredContours, currentPipeline.target_intersection, currentPipeline.target_group);
                if (GroupedContours.size() > 0) {
                    var finalRect = visionProcess.SortTargetsToOne(GroupedContours, currentPipeline.sort_mode);
                    pipelineResult.RawPoint = finalRect;
                    pipelineResult.IsValid = true;
                    if (!currentPipeline.is_calibrated) {
                        pipelineResult.CalibratedX = camVals.CenterX;
                        pipelineResult.CalibratedY = camVals.CenterY;
                    } else {
                        pipelineResult.CalibratedX = (finalRect.center.y - currentPipeline.B) / currentPipeline.M;
                        pipelineResult.CalibratedY = finalRect.center.x * currentPipeline.M + currentPipeline.B;
                        pipelineResult.Pitch = camVals.CalculatePitch(finalRect.center.y, pipelineResult.CalibratedY);
                        pipelineResult.Yaw = camVals.CalculateYaw(finalRect.center.x, pipelineResult.CalibratedX);
                    }
                    // TODO Send pitch yaw distance and Raw Point using websockets to client for calib calc
                    drawContour(outputImage, finalRect);
                }
            }
        }

        return pipelineResult;
    }

    @Override
    public void run() {
        // processing time tracking
        long startTime;
        double processTimeMs;
        double fps = 0;

        while (!Thread.interrupted()) {
            FoundContours.clear();
            FilteredContours.clear();
            GroupedContours.clear();

            currentPipeline = camera.getCurrentPipeline();
            // start fps counter right before grabbing input frame
            startTime = System.nanoTime();
            TimeStamp = camera.grabFrame(cameraInputMat);
            if (cameraInputMat.cols() == 0 && cameraInputMat.rows() == 0) {
                continue;
            }

            // get vision data
            var pipelineResult = runVisionProcess(cameraInputMat, streamOutputMat);
            updateNetworkTables(pipelineResult);
            if (cameraName.equals(SettingsManager.GeneralSettings.curr_camera) && pipelineResult.IsValid) {
                HashMap<String,Object> WebSend = new HashMap<>();
                HashMap<String,Object> point = new HashMap<>();
                List<Double> center = new ArrayList<>();
                center.add(pipelineResult.RawPoint.center.x);
                center.add(pipelineResult.RawPoint.center.y);
                point.put("pitch", pipelineResult.Pitch);
                point.put("yaw", pipelineResult.Yaw);
                point.put("fps", fps);
                WebSend.put("point", point);
                WebSend.put("raw_point", center);
                Server.broadcastMessage(WebSend);
            }
            camera.putFrame(streamOutputMat);
            // calculate FPS after publishing output frame
            processTimeMs = (System.nanoTime() - startTime) * 1e-6;
            fps = 1000 / processTimeMs;
            System.out.printf("%s - Process time: %fms, FPS: %.2f, FoundContours: %d, FilteredContours: %d, GroupedContours: %d\n", cameraName,processTimeMs, fps, FoundContours.size(), FilteredContours.size(), GroupedContours.size());

            cameraInputMat.release();
            hsvThreshMat.release();

            memManager.run(true);
        }
    }
}