Add python samples and code

SamCarlberg · SamCarlberg · commit 5d01515cf6d9 · 2016-12-15T10:51:12.000-05:00
diff --git a/python/README.md b/python/README.md
@@ -0,0 +1,41 @@
+# Python samples and code
+
+The `files` directory contains skeleton code that can be modified to run a GRIP-generated pipeline.
+
+The `samples` directory has sample programs that can be run directly. Just clone this repository and run the main python files.
+
+## Using the samples
+
+The samples use Python 3 and require three libraries:
+
+1. OpenCV
+2. numpy (rquired by OpenCV)
+3. pynetworktables
+
+These can all be installed with pip3.
+
+```bash
+$ pip3 install numpy opencv-python pynetworktables
+```
+
+(if your default python version is at least 3+, the normal `pip` command should work)
+
+
+Then `cd` into the directory of the sample you want to run and use the `python3` command to run the sample, eg
+
+```bash
+$ cd samples/frc_find_red_areas
+$ python3 frc_find_red_areas.py
+```
+
+(Again, if your default python version is at least 3+, just the `python` command should work)
+
+
+## Binaries
+
+Since the Raspberry Pi 3 is a cheap ($35) and easy to use co-processor, we've created a Python 3 package for OpenCV for it and is available on the releases page. To install it, you'll need to have the python `wheel` tool
+
+```bash
+$ pip3 install wheel
+$ python3 -m wheel install opencv_python_rpi3-...whl
+```
diff --git a/python/files/GripRunner.py b/python/files/GripRunner.py
@@ -0,0 +1,42 @@
+#!/usr/bin/python3
+
+"""
+Simple skeleton program for running an OpenCV pipeline generated by GRIP and using NetworkTables to send data.
+
+Users need to:
+
+1. Import the generated GRIP pipeline, which should be generated in the same directory as this file.
+2. Set the network table server IP. This is usually the robots address (roborio-TEAM-frc.local) or localhost
+3. Handle putting the generated code into NetworkTables
+"""
+
+import cv2
+from networktables import NetworkTable
+from grip import GripPipeline  # TODO change the default module and class, if needed
+
+
+def extra_processing(pipeline: GripPipeline):
+    """
+    Performs extra processing on the pipeline's outputs and publishes data to NetworkTables.
+    :param pipeline: the pipeline that just processed an image
+    :return: None
+    """
+    # TODO: Users need to implement this.
+    # Useful for converting OpenCV objects (e.g. contours) to something NetworkTables can understand.
+    pass
+
+
+def main():
+    NetworkTable.setTeam('0000')  # TODO set your team number
+    NetworkTable.initialize()
+    cap = cv2.VideoCapture(0)
+    pipeline = GripPipeline()
+    while True:
+        ret, frame = cap.read()
+        if ret:
+            pipeline.process(frame)  # TODO add extra parameters if the pipeline takes more than just a single image
+            extra_processing(pipeline)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/python/samples/frc_find_red_areas/frc_find_red_areas.py b/python/samples/frc_find_red_areas/frc_find_red_areas.py
@@ -0,0 +1,62 @@
+#!/usr/bin/python3
+
+"""
+Sample program that uses a generated GRIP pipeline to detect red areas in an image and publish them to NetworkTables.
+"""
+
+import cv2
+from ntcore import NetworkTables
+from grip import GripPipeline
+
+
+def extra_processing(pipeline):
+    """
+    Performs extra processing on the pipeline's outputs and publishes data to NetworkTables.
+    :param pipeline: the pipeline that just processed an image
+    :return: None
+    """
+    center_x_positions = []
+    center_y_positions = []
+    widths = []
+    heights = []
+
+    # Find the bounding boxes of the contours to get x, y, width, and height
+    for contour in pipeline.filter_contours_output:
+        x, y, w, h = cv2.boundingRect(contour)
+        center_x_positions.append(x + w / 2)  # X and Y are coordinates of the top-left corner of the bounding box
+        center_y_positions.append(y + h / w)
+        widths.append(w)
+        heights.append(y)
+
+    # Publish to the '/vision/red_areas' network table
+    table = NetworkTables.getTable('/vision/red_areas')
+    table.putNumberArray('x', center_x_positions)
+    table.putNumberArray('y', center_y_positions)
+    table.putNumberArray('width', widths)
+    table.putNumberArray('height', heights)
+
+
+def main():
+    print('Initializing NetworkTables')
+    NetworkTables.setClientMode()
+    NetworkTables.setIPAddress('localhost')
+    NetworkTables.initialize()
+
+    print('Creating video capture')
+    cap = cv2.VideoCapture(0)
+
+    print('Creating pipeline')
+    pipeline = GripPipeline()
+
+    print('Running pipeline')
+    while cap.isOpened():
+        have_frame, frame = cap.read()
+        if have_frame:
+            pipeline.process(frame)
+            extra_processing(pipeline)
+
+    print('Capture closed')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/python/samples/frc_find_red_areas/grip.py b/python/samples/frc_find_red_areas/grip.py
@@ -0,0 +1,160 @@
+import cv2
+import numpy
+import math
+from enum import Enum
+
+class GripPipeline:
+    """
+    An OpenCV pipeline generated by GRIP.
+    """
+    
+    def __init__(self):
+        """initializes all values to presets or None if need to be set
+        """
+
+        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.__rgb_threshold_input = self.resize_image_output
+        self.__rgb_threshold_red = [91.72661870503596, 255.0]
+        self.__rgb_threshold_green = [0.0, 53.684210526315795]
+        self.__rgb_threshold_blue = [0.0, 73.16638370118847]
+
+        self.rgb_threshold_output = None
+
+        self.__find_contours_input = self.rgb_threshold_output
+        self.__find_contours_external_only = True
+
+        self.find_contours_output = None
+
+        self.__filter_contours_contours = self.find_contours_output
+        self.__filter_contours_min_area = 20.0
+        self.__filter_contours_min_perimeter = 0
+        self.__filter_contours_min_width = 0
+        self.__filter_contours_max_width = 1000
+        self.__filter_contours_min_height = 0
+        self.__filter_contours_max_height = 1000
+        self.__filter_contours_solidity = [0, 100]
+        self.__filter_contours_max_vertices = 1000000
+        self.__filter_contours_min_vertices = 0
+        self.__filter_contours_min_ratio = 0
+        self.__filter_contours_max_ratio = 1000
+
+        self.filter_contours_output = None
+
+
+    def process(self, source0):
+        """
+        Runs the pipeline and sets all outputs to new values.
+        """
+        # Step Resize_Image0:
+        self.__resize_image_input = source0
+        (self.resize_image_output) = self.__resize_image(self.__resize_image_input, self.__resize_image_width, self.__resize_image_height, self.__resize_image_interpolation)
+
+        # Step RGB_Threshold0:
+        self.__rgb_threshold_input = self.resize_image_output
+        (self.rgb_threshold_output) = self.__rgb_threshold(self.__rgb_threshold_input, self.__rgb_threshold_red, self.__rgb_threshold_green, self.__rgb_threshold_blue)
+
+        # Step Find_Contours0:
+        self.__find_contours_input = self.rgb_threshold_output
+        (self.find_contours_output) = self.__find_contours(self.__find_contours_input, self.__find_contours_external_only)
+
+        # Step Filter_Contours0:
+        self.__filter_contours_contours = self.find_contours_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 __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 __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 __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 __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
+
+
+
diff --git a/sample-pipelines/find_red_areas.grip b/sample-pipelines/find_red_areas.grip
