frc2554_vision_final.py

#!/usr/bin/env python3

# ---------------------------------------- #
#             Begin GRIP Pipeline          #
# ---------------------------------------- #


import cv2
cv2.setUseOptimized(True)
import numpy
import math
from enum import Enum

class VisionPipeline:
    """
    An OpenCV pipeline generated by GRIP.
    """
    
    def __init__(self):
        """initializes all values to presets or None if need to be set
        """

        self.__rgb_threshold_red = [0.0, 144.92648866498453]
        self.__rgb_threshold_green = [145.14312278990366, 255.0]
        self.__rgb_threshold_blue = [0.0, 144.49355199942204]

        self.rgb_threshold_output = None

        self.__resize_image_input = self.rgb_threshold_output
        self.__resize_image_width = 320.0
        self.__resize_image_height = 240.0
        self.__resize_image_interpolation = cv2.INTER_CUBIC

        self.resize_image_output = None

        self.__find_contours_input = self.resize_image_output
        self.__find_contours_external_only = True

        self.find_contours_output = None

        self.__convex_hulls_contours = self.find_contours_output

        self.convex_hulls_output = None

        self.__filter_contours_contours = self.convex_hulls_output
        self.__filter_contours_min_area = 20.0
        self.__filter_contours_min_perimeter = 0.0
        self.__filter_contours_min_width = 0.0
        self.__filter_contours_max_width = 1000.0
        self.__filter_contours_min_height = 0.0
        self.__filter_contours_max_height = 1000.0
        self.__filter_contours_solidity = [0, 100]
        self.__filter_contours_max_vertices = 1000000.0
        self.__filter_contours_min_vertices = 0.0
        self.__filter_contours_min_ratio = 0.0
        self.__filter_contours_max_ratio = 1000.0

        self.filter_contours_output = None


    def process(self, source0):
        """
        Runs the pipeline and sets all outputs to new values.
        """
        # Step RGB_Threshold0:
        self.__rgb_threshold_input = source0
        (self.rgb_threshold_output) = self.__rgb_threshold(self.__rgb_threshold_input, self.__rgb_threshold_red, self.__rgb_threshold_green, self.__rgb_threshold_blue)

        # Step Resize_Image0:
        self.__resize_image_input = self.rgb_threshold_output
        (self.resize_image_output) = self.__resize_image(self.__resize_image_input, self.__resize_image_width, self.__resize_image_height, self.__resize_image_interpolation)

        # Step Find_Contours0:
        self.__find_contours_input = self.resize_image_output
        (self.find_contours_output) = self.__find_contours(self.__find_contours_input, self.__find_contours_external_only)

        # Step Convex_Hulls0:
        self.__convex_hulls_contours = self.find_contours_output
        (self.convex_hulls_output) = self.__convex_hulls(self.__convex_hulls_contours)

        # Step Filter_Contours0:
        self.__filter_contours_contours = self.convex_hulls_output
        (self.filter_contours_output) = self.__filter_contours(self.__filter_contours_contours, self.__filter_contours_min_area, self.__filter_contours_min_perimeter, self.__filter_contours_min_width, self.__filter_contours_max_width, self.__filter_contours_min_height, self.__filter_contours_max_height, self.__filter_contours_solidity, self.__filter_contours_max_vertices, self.__filter_contours_min_vertices, self.__filter_contours_min_ratio, self.__filter_contours_max_ratio)


    @staticmethod
    def __rgb_threshold(input, red, green, blue):
        """Segment an image based on color ranges.
        Args:
            input: A BGR numpy.ndarray.
            red: A list of two numbers the are the min and max red.
            green: A list of two numbers the are the min and max green.
            blue: A list of two numbers the are the min and max blue.
        Returns:
            A black and white numpy.ndarray.
        """
        out = cv2.cvtColor(input, cv2.COLOR_BGR2RGB)
        return cv2.inRange(out, (red[0], green[0], blue[0]),  (red[1], green[1], blue[1]))

    @staticmethod
    def __resize_image(input, width, height, interpolation):
        """Scales and image to an exact size.
        Args:
            input: A numpy.ndarray.
            Width: The desired width in pixels.
            Height: The desired height in pixels.
            interpolation: Opencv enum for the type fo interpolation.
        Returns:
            A numpy.ndarray of the new size.
        """
        return cv2.resize(input, ((int)(width), (int)(height)), 0, 0, interpolation)

    @staticmethod
    def __find_contours(input, external_only):
        """Sets the values of pixels in a binary image to their distance to the nearest black pixel.
        Args:
            input: A numpy.ndarray.
            external_only: A boolean. If true only external contours are found.
        Return:
            A list of numpy.ndarray where each one represents a contour.
        """
        if(external_only):
            mode = cv2.RETR_EXTERNAL
        else:
            mode = cv2.RETR_LIST
        method = cv2.CHAIN_APPROX_SIMPLE
        im2, contours, hierarchy =cv2.findContours(input, mode=mode, method=method)
        return contours

    @staticmethod
    def __convex_hulls(input_contours):
        """Computes the convex hulls of contours.
        Args:
            input_contours: A list of numpy.ndarray that each represent a contour.
        Returns:
            A list of numpy.ndarray that each represent a contour.
        """
        output = []
        for contour in input_contours:
            output.append(cv2.convexHull(contour))
        return output

    @staticmethod
    def __filter_contours(input_contours, min_area, min_perimeter, min_width, max_width,
                        min_height, max_height, solidity, max_vertex_count, min_vertex_count,
                        min_ratio, max_ratio):
        """Filters out contours that do not meet certain criteria.
        Args:
            input_contours: Contours as a list of numpy.ndarray.
            min_area: The minimum area of a contour that will be kept.
            min_perimeter: The minimum perimeter of a contour that will be kept.
            min_width: Minimum width of a contour.
            max_width: MaxWidth maximum width.
            min_height: Minimum height.
            max_height: Maximimum height.
            solidity: The minimum and maximum solidity of a contour.
            min_vertex_count: Minimum vertex Count of the contours.
            max_vertex_count: Maximum vertex Count.
            min_ratio: Minimum ratio of width to height.
            max_ratio: Maximum ratio of width to height.
        Returns:
            Contours as a list of numpy.ndarray.
        """
        output = []
        for contour in input_contours:
            x,y,w,h = cv2.boundingRect(contour)
            if (w < min_width or w > max_width):
                continue
            if (h < min_height or h > max_height):
                continue
            area = cv2.contourArea(contour)
            if (area < min_area):
                continue
            if (cv2.arcLength(contour, True) < min_perimeter):
                continue
            hull = cv2.convexHull(contour)
            solid = 100 * area / cv2.contourArea(hull)
            if (solid < solidity[0] or solid > solidity[1]):
                continue
            if (len(contour) < min_vertex_count or len(contour) > max_vertex_count):
                continue
            ratio = (float)(w) / h
            if (ratio < min_ratio or ratio > max_ratio):
                continue
            output.append(contour)
        return output

# ---------------------------------------- #
#             End GRIP Pipeline            #
# ---------------------------------------- #

# ---------------------------------------- #
#             Begin FRC Template           #
# ---------------------------------------- #

# ----------------------------------------------------------------------------
# Copyright (c) 2018 FIRST. All Rights Reserved.
# Open Source Software - may be modified and shared by FRC teams. The code
# must be accompanied by the FIRST BSD license file in the root directory of
# the project.
# ----------------------------------------------------------------------------

import json
import time
import sys

from cscore import CameraServer, VideoSource, UsbCamera, MjpegServer
from networktables import NetworkTablesInstance

#   JSON format:
#   {
#       "team": <team number>,
#       "ntmode": <"client" or "server", "client" if unspecified>
#       "cameras": [
#           {
#               "name": <camera name>
#               "path": <path, e.g. "/dev/video0">
#               "pixel format": <"MJPEG", "YUYV", etc>   // optional
#               "width": <video mode width>              // optional
#               "height": <video mode height>            // optional
#               "fps": <video mode fps>                  // optional
#               "brightness": <percentage brightness>    // optional
#               "white balance": <"auto", "hold", value> // optional
#               "exposure": <"auto", "hold", value>      // optional
#               "properties": [                          // optional
#                   {
#                       "name": <property name>
#                       "value": <property value>
#                   }
#               ],
#               "stream": {                              // optional
#                   "properties": [
#                       {
#                           "name": <stream property name>
#                           "value": <stream property value>
#                       }
#                   ]
#               }
#           }
#       ]
#   }

configFile = "/boot/frc.json"
config_json = '{ "fps": 60, "height": 480, "pixel format": "mjpeg", "properties": [ { "name": "connect_verbose", "value": 1 }, { "name": "raw_brightness", "value": 135 }, { "name": "brightness", "value": 53 }, { "name": "raw_contrast", "value": 81 }, { "name": "contrast", "value": 32 }, { "name": "raw_saturation", "value": 132 }, { "name": "saturation", "value": 52 }, { "name": "white_balance_temperature_auto", "value": false }, { "name": "raw_gain", "value": 40 }, { "name": "gain", "value": 16 }, { "name": "power_line_frequency", "value": 2 }, { "name": "white_balance_temperature", "value": 6500 }, { "name": "raw_sharpness", "value": 20 }, { "name": "sharpness", "value": 8 }, { "name": "backlight_compensation", "value": 1 }, { "name": "exposure_auto", "value": 1 }, { "name": "raw_exposure_absolute", "value": 23 }, { "name": "exposure_absolute", "value": 1 }, { "name": "exposure_auto_priority", "value": true }, { "name": "pan_absolute", "value": 0 }, { "name": "tilt_absolute", "value": 0 }, { "name": "focus_absolute", "value": 51 }, { "name": "focus_auto", "value": true }, { "name": "zoom_absolute", "value": 1 } ], "width": 640 }'


class CameraConfig:
    pass


team = 2554
server = False
cameraConfigs = []

"""Report parse error."""


def parseError(str):
    print("config error in '" + configFile + "': " + str, file=sys.stderr)


"""Read single camera configuration."""


def readCameraConfig(config):
    cam = CameraConfig()

    # name
    try:
        cam.name = config["name"]
    except KeyError:
        parseError("could not read camera name")
        return False

    # path
    try:
        cam.path = config["path"]
    except KeyError:
        parseError("camera '{}': could not read path".format(cam.name))
        return False

    # stream properties
    cam.streamConfig = config.get("stream")

    cam.config = config

    cameraConfigs.append(cam)
    return True


"""Read configuration file."""


def readConfig():
    global team
    global server

    # parse file
    try:
        with open(configFile, "rt") as f:
            j = json.load(f)
    except OSError as err:
        print("could not open '{}': {}".format(configFile, err), file=sys.stderr)
        return False

    # top level must be an object
    if not isinstance(j, dict):
        parseError("must be JSON object")
        return False

    # team number
    try:
        team = j["team"]
    except KeyError:
        parseError("could not read team number")
        return False

    # ntmode (optional)
    if "ntmode" in j:
        str = j["ntmode"]
        if str.lower() == "client":
            server = False
        elif str.lower() == "server":
            server = True
        else:
            parseError("could not understand ntmode value '{}'".format(str))

    # cameras
    try:
        cameras = j["cameras"]
    except KeyError:
        parseError("could not read cameras")
        return False
    for camera in cameras:
        if not readCameraConfig(camera):
            return False

    return True


"""Start running the camera."""


def startCamera(config):
    print("Starting camera '{}' on {}".format(config.name, config.path))
    inst = CameraServer.getInstance()
    camera = UsbCamera(config.name, config.path)

    # camera.setConfigJson(json.dumps(config.config))
    camera.setConfigJson(config_json)
    camera.setConnectionStrategy(VideoSource.ConnectionStrategy.kKeepOpen)

    inst.addCamera(camera)

    return inst, camera


# ---------------------------------------- #
#             End FRC Template             #
# ---------------------------------------- #

# ---------------------------------------- #
#             Begin Our Code               #
# ---------------------------------------- #

import cv2
cv2.setUseOptimized(True)
from math import tan, sqrt
import numpy as np

IMAGE_WIDTH = 320
IMAGE_HEIGHT = 240

HFOV = 65.8725303703

DEG_PER_PIXEL = HFOV / IMAGE_WIDTH

CENTER_WIDTH_PIXEL = (IMAGE_WIDTH - 1) // 2
CENTER_HEIGHT_PIXEL = (IMAGE_HEIGHT - 1) // 2


def getContourAngle(contour):
    rect = cv2.minAreaRect(contour)
    angle = rect[-1]

    return angle


def angleToTarget(img, contours):
    new_image = img

    contour_diff = -500
    c1a = -250
    c2a = -250
    angle = -420
    center1 = (21, 69)
    center2 = (420, 666)
    targetCenter = (999, 999)
    pixelDiff = -6969
    targetExists = False

    imgCenter = (CENTER_WIDTH_PIXEL, CENTER_HEIGHT_PIXEL)
    cv2.circle(
        img=new_image, center=(imgCenter), radius=3, color=(255, 0, 0), thickness=-1
    )

    if len(contours) >= 2:
        contours = list(sorted(contours, key=cv2.contourArea))[::-1]
        cntAngles = [getContourAngle(i) for i in contours]

        finalCnts = []
        baseAngle = -69

        for idx, i in enumerate(cntAngles):
            if idx == 0:
                finalCnts.append(contours[idx])
                baseAngle = i
            else:
                diff = abs((i % 360) - (baseAngle % 360))
                c1a = i % 360
                c2a = baseAngle % 360
                contour_diff = diff
                if diff > 55 and diff < 80:
                    finalCnts.append(contours[idx])
                    break
        if not len(finalCnts) < 2:
            finalCnts = list(sorted(finalCnts, key=cv2.contourArea))[::-1]
            cnt1 = finalCnts[0]
            cnt2 = finalCnts[1]

            cv2.drawContours(new_image, finalCnts, -1, color=(255, 0, 0), thickness=2)

            M1 = cv2.moments(cnt1)
            M2 = cv2.moments(cnt2)

            center1 = (int(M1["m10"] / M1["m00"]), int(M1["m01"] / M1["m00"]))
            center2 = (int(M2["m10"] / M2["m00"]), int(M2["m01"] / M2["m00"]))

            # cv2.circle(
            #     img=new_image, center=center1, radius=3, color=(0, 0, 255), thickness=-1
            # )
            # cv2.circle(
            #     img=new_image, center=center2, radius=3, color=(0, 0, 255), thickness=-1
            # )

            # Draw the midpoint of both of these contours
            targetCenter = (
                int((center1[0] + center2[0]) / 2),
                int((center1[1] + center2[1]) / 2),
            )
            cv2.circle(
                img=new_image,
                center=targetCenter,
                radius=3,
                color=(0, 0, 255),
                thickness=-1,
            )

            angle = (targetCenter[0] - CENTER_WIDTH_PIXEL) * DEG_PER_PIXEL
            targetExists = True
            cv2.putText(
                new_image,
                str(round(angle, 2)) + " deg",
                (0, 25),
                cv2.FONT_HERSHEY_SIMPLEX,
                1,
                color=(0, 255, 255),
                thickness=2,
            )

            cv2.line(new_image, targetCenter, imgCenter, (255, 0, 0), 2)
            pixelDiff = targetCenter[0] - imgCenter[0]
    
    shuffleboard_data = {
        "target_exists": targetExists,
        "center1": center1,
        "center2": center2,
        "midpoint": targetCenter,
        "pixel_diff": pixelDiff,
        "yaw_angle": angle,
        "contour_diff": contour_diff,
        "c1a": c1a,
        "c2a": c2a
    }
    return new_image, shuffleboard_data

from threading import Thread

class ThreadedVision:
    def __init__(self, frame):
        self.grip = VisionPipeline()
        self.running = True
        self.frame = frame
        self.output = None
    def start(self):
        Thread(target=self.run, args=()).start()
        return self
    def run(self):
        while self.running:
            frame = self.frame.copy()
            self.grip.process(frame)
            frame = cv2.resize(frame, (320, 240), 0, 0, cv2.INTER_CUBIC)
            self.output = angleToTarget(frame, self.grip.filter_contours_output)

image_width = 640
image_height = 480

class ThreadedInput:
    def __init__(self, cvSink):
        self.img = np.zeros(shape=(image_height, image_width, 3), dtype=np.uint8)
        self.cvSink = cvSink
        self.timestamp = 0
    def start(self):
        Thread(target=self.run, args=()).start()
        return self
    def run(self):
        while True:
            self.timestamp, self.img = self.cvSink.grabFrame(self.img) 
            pass

def main():
    global configFile

    if len(sys.argv) >= 2:
        configFile = sys.argv[1]

    if not readConfig():
        print("Unable to read config file!")
        sys.exit(1)

    # start cameras
    streams = []

    print("Initialized vision stuff")

    for cameraConfig in cameraConfigs:
        # cameras.append(startCamera(cameraConfig))
        cs, cameraCapture = startCamera(cameraConfig)
        streams.append(cs)

    # First camera is server
    cameraServer = streams[0]

    # Set up a CV Sink to capture video
    cvSink = cameraServer.getVideo()

    # CvSource
    outputStream = cameraServer.putVideo("stream", image_width, image_height)

    img = np.zeros(shape=(image_height, image_width, 3), dtype=np.uint8)

    # Networktables
    ninst = NetworkTablesInstance.getDefault()
    if server:
        print("Setting up NetworkTables server")
        ninst.startServer()
    else:
        print("Setting up NetworkTables client for team {}".format(team))
        ninst.startClientTeam(team)

    network_table = ninst.getTable("Shuffleboard").getSubTable("Vision")
    network_table.getEntry("connected").setValue(True)

    imgetter = ThreadedInput(cvSink).start()
    timestamp, img = imgetter.timestamp, imgetter.img
    vis = ThreadedVision(img).start()
    
    time.sleep(0.1)

    num_frames = 0

    while True:
        start = time.time()
        num_frames += 1
        start = time.time()
        timestamp, img = imgetter.timestamp, imgetter.img
        vis.frame = img

        if timestamp == 0:
            outputStream.notifyError(cvSink.getError())
            continue

        new_image, shuffleboard_data = vis.output

        for name, data in shuffleboard_data.items():
           network_table.getEntry(name).setValue(data)

        new_image = cv2.resize(new_image, (160, 120))

        outputStream.putFrame(new_image)

        fps = 1 / (time.time() - start)
        if num_frames % 1000 == 0:
            print(fps)
            num_frames = 0
        time.sleep(max(1.0 / 30.0 - (time.time() - start), 0))

if __name__ == "__main__":
    main()

# ---------------------------------------- #
#             End Our Code                 #
# ---------------------------------------- #