vision.py 文件源码

def find_target(img, lower=np.array([110//2, 10*255//100, 15*255//100]), upper=np.array([180//2, 100*255//100, 100*255//100]), area_threshold=0.025 ** 2):
    """Given an image and thresholds, find the centre of mass of the target.
    All arguments must be np.arrays, except for area_threshold, and lower and upper must be a 3-array.
    """

    #Converting from RGB to HSV.
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    #Making the mask.
    mask = cv2.inRange(hsv, lower, upper)

    #Combining the mask with the frame
    res = cv2.bitwise_and(img, img, mask=mask)

    height, width = mask.shape
    # Get the information for the contours
    _, contours, __ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # sort the contours into a list
    areas = []
    for idx, contour in enumerate(contours):
        area = cv2.contourArea(contour)
        if area/(height*width) > area_threshold:
            heapq.heappush(areas, (cv2.contourArea(contour), idx))

    areas = heapq.nlargest(2, areas)
    areas_x = []
    x_coord = 0
    for _, idx in areas:
        contour = contours[idx]
        moments = cv2.moments(contour)
        x_coord += moments['m10']/moments['m00'] / len(areas)
        areas_x.append(moments['m10']/moments['m00'])
        cv2.drawContours(res, (contour, ), -1, (255, 0, 0), 1)
    if len(areas) > 0:
        cv2.line(res, (int(x_coord), 60), (int(x_coord), 180), (255,255,0), thickness=2, lineType=8, shift=0)
    target_sep = 0
    if len(areas_x) > 1:
        # target sep returned as a % of image width, not in vision coordinates
        target_sep = abs(areas_x[0]-areas_x[1]) / width
    pos = 2 * x_coord / width - 1
    return pos, res, len(areas), target_sep

# Allow easy testing of captured sample images