Vision AI using OpenCV

Question

I know this may sound dumb, but last time, I posted C++ code, similar to this one:

This is just something to get me closer to actually doing a neural network. This code is supposed to learn colors from many many data images, and then recognize color using an algorithm that I made. Eventually, I want the program to make its own algorithm.

That was my prototype to the Python version. Why you ask? Because I am better in C++ so I prototyped it in C++, but I want the final version to be in Python so bringing it onto a Raspberry Pi would be easier.

The lower HSV target getting part is quite sketchy so hopefully someone can make it better. My goal is to make the fastest PYTHON program possible.

Code and training data is on GitHub. Consider contributing directly to the repository.

import numpy as np
import cv2 as cv
import os

SCALAR_SIZE = 3

COLOR_NAME = 0
COLOR_BGR = 1
COLOR_DIFFERENCE = 2
COLOR_ACCURACY = 3

DIR_TRAIN_DATA = "train_data"
DIR_TEST_DATA = "test_data"
DIR_SAVED_DATA = "saved_data"

DIR_NAME = "name"
FILE_NAME = "name.txt"

DIR_IMAGE = "images"

FILE_SAVED_HSV = "hsv_values.txt"

def nothing(something):
    pass

def get_bgr_difference(bgr):
    return [bgr[0] - bgr[1], bgr[1] - bgr[2], bgr[2] - bgr[0]]

def get_color(image, colors):
    difference = get_bgr_difference(np.average(np.average(image, axis=0), axis=0))
    accuracy = []
    for color in colors:
        accuracy.append(1 - (np.average(abs(np.subtract((color[COLOR_DIFFERENCE], difference)))) / 255))
    color_accuracy = max(accuracy)
    color_match = colors[accuracy.index(color_accuracy)]
    color_match[COLOR_ACCURACY] = color_accuracy
    return color_match

def get_trained_colors():
    color = []
    for train_data_folder in os.walk(DIR_TRAIN_DATA):
        for color_name in train_data_folder[1]:
            if color_name != DIR_NAME and color_name != DIR_IMAGE:
                location = DIR_TRAIN_DATA + '/' + color_name + '/'
                bgr = []
                for color_images in os.walk(location + '/' + DIR_IMAGE):
                    for image_file in color_images[2]:
                        bgr.append(
                            np.average(np.average(cv.imread(location + '/' + DIR_IMAGE + '/' + image_file, cv.IMREAD_COLOR), axis=0),
                                       axis=0))
                bgr = np.average(bgr, axis=0)
                color.append([open(location + DIR_NAME + '/' + FILE_NAME).read(), bgr, get_bgr_difference(bgr), 1])
    return color

def get_position_in_list(myList, v):
    for i, x in enumerate(myList):
        if v in x:
            return i, x.index(v)

def get_target_image_bgr(colors, image, target_color_name, tolerance):
    color = colors[get_position_in_list(colors, target_color_name)[0]][COLOR_BGR]
    image = cv.inRange(cv.blur(image, (15, 15)), np.subtract(color, tolerance), np.add(color, tolerance))
    return image

def get_target_image_hsv(image, tolerance):
    cv.threshold(image, 0, 255, cv.THRESH_BINARY_INV)
    cv.cvtColor(image, cv.COLOR_BGR2HSV)
    image = cv.inRange(cv.cvtColor(image, cv.COLOR_BGR2HSV), tolerance[0], tolerance[1])
    kernel = np.ones((10, 10), np.uint8)
    image = cv.dilate(image, kernel, iterations=1)
    cv.imshow("hsv", image)
    return image

def get_target_coordinate(image):
    x = cv.findContours(image, cv.RETR_TREE, cv.CHAIN_APPROX_SIMPLE)[1]
    y = []
    cnt = []
    for i in x:
        y.append(np.average(i, axis=0))
    cnt = np.average(y, axis=0)[0]
    return cnt

def draw_target(image, coordinates):
    cv.circle(image, (int(coordinates[0]), int(coordinates[1])), 5, (255, 0, 255), -1)
    cv.imshow("image", image)
    cv.waitKey(1)

def draw_trackbar_hsv():
    window_name = 'tracker'
    cv.namedWindow(window_name)
    for i in ['h', 's', 'v']:
        for j in range(2):
            cv.createTrackbar(i + str(j), 'tracker', 0, 255, nothing)

def get_trackbar():
    hsv = np.array([[0] * 3] * 2)
    hsv[0][0] = cv.getTrackbarPos('h0', 'tracker')
    hsv[1][0] = cv.getTrackbarPos('h1', 'tracker')
    hsv[0][1] = cv.getTrackbarPos('s0', 'tracker')
    hsv[1][1] = cv.getTrackbarPos('s1', 'tracker')
    hsv[0][2] = cv.getTrackbarPos('v0', 'tracker')
    hsv[1][2] = cv.getTrackbarPos('v1', 'tracker')
    return hsv

def set_trackbar():
    text_file = open(DIR_SAVED_DATA + '/' + FILE_SAVED_HSV, "r")
    hsv = text_file.read().split(',')
    text_file.close()
    hsv.remove('')
    hsv = [int(i) for i in hsv]
    cv.setTrackbarPos('h0', 'tracker', hsv[0])
    cv.setTrackbarPos('h1', 'tracker', hsv[3])
    cv.setTrackbarPos('s0', 'tracker', hsv[1])
    cv.setTrackbarPos('s1', 'tracker', hsv[4])
    cv.setTrackbarPos('v0', 'tracker', hsv[2])
    cv.setTrackbarPos('v1', 'tracker', hsv[5])

def save_trackbar_hsv(hsv):
    name = DIR_SAVED_DATA + '/' + FILE_SAVED_HSV
    open(name, "w").close()
    text_file = open(name, "w")
    for i in hsv:
        for j in i:
            text_file.write(str(j))
            text_file.write(',')
    text_file.close()

image = cv.imread(DIR_TEST_DATA + '/' + "boiler3.jpg", cv.IMREAD_COLOR)
cv.namedWindow("image")
cv.namedWindow("hsv")
draw_trackbar_hsv()
set_trackbar()
while True:
    hsv_val = get_trackbar()
    hsv_image = get_target_image_hsv(image, hsv_val)
    coordinate = get_target_coordinate(hsv_image)
    save_trackbar_hsv(hsv_val)
    draw_target(image, coordinate)
cv.destroyAllWindows()

Answer 1

Use a dictionary

The following function looks like it might get slow if it's called often:

def get_position_in_list(myList, v):
    for i, x in enumerate(myList):
        if v in x:
            return i, x.index(v)

If colors is a large list of lists (say 100 by 100), then searching through it one by one will give you a significant performance hit.

I'm not sure how often you plan to run get_target_image_bgr, but it's probably going to be often, especially if it'll be in a while True loop, similar to get_target_image_hsv.

In that case, when you create the 'colors' argument to get_target_image_bgr, use a dictionary instead of a list of lists:

Initializing the dictionary:

colors = {}
for i,_ in enumerate(myList):
    for j, v in enumerate(myList[i]):
        colors[v] = (i, j)

Looking up values in the dictionary:

colors[v]

This will return the same co-ordinates as your get_position_in_list function, but each lookup takes \$\mathcal{O}(1)\$ instead of \$\mathcal{O}(n^2)\$.

(Note: I have assumed v is unique. If not, the dictionary as I have it will return the last match, while your get_position_in_list function returns the first match.)

Reduce writes to file

With new SSDs getting up to 4GB/s write, this old axiom may not be as important as it once was. Still, I think you should change up your save_trackbar_hsv, you will probably notice a big performance difference.

Currently, you are writing a two-dimensional array to file with every iteration of the while True loop. Not only that, but you are also making a separate call to write for every data value.

Your save_trackbar_hsv should make only one call to write, writing all data values in one go, with appropriate delimiters:

def save_trackbar_hsv(hsv):
    name = DIR_SAVED_DATA + '/' + FILE_SAVED_HSV
    open(name, "w").close()
    text_file = open(name, "w")
    between_values_delimiter = " "
    between_lines_delimiter = ","
    file_data = []

    for i in hsv:
        file_line = []
        for j in i:
            file_line.append(str(j))
        file_data.append(between_values_delimiter.join(file_line))
    text_file.write(between_lines_delimiter.join(file_data))

You may also wish to store file data for multiple calls to save_trackbar_hsv in a list, and only write on perhaps each 100th call to the funcion.

On the positive side: good use of numpy, opencv and os.walk.

Lastly, sometimes you may hear the attitude that any program written in C++ is by default faster than any program written in Python. That is certainly not true.

Looks OK overall, and a great effort! Congrats for branching out to Python.