tinyPnP/pnpMachineFlask/frame.py

#!/usr/bin/python3
# -*- coding: utf-8 -*-

import cv2
import numpy as np
import math
import pytesseract

def frame_grid(frame, camParams):
    h, w = frame.shape[:2]
    cx = int(w/2)
    cy = int(h/2)
    color = (0, 0, 255)

    if camParams['grid']:
        frame = cv2.line(frame, (cx, 0), (cx, h), color, thickness=1, lineType=8)
        frame = cv2.line(frame, (0, cy), (w, cy), color, thickness=1, lineType=8)

    if camParams['marker']:
        ppmm_h = 10*camParams['ppmm']/2
        frame = cv2.rectangle(frame, (int(cx-ppmm_h), int(cy-ppmm_h)), (int(cx+ppmm_h), int(cy+ppmm_h)), color, thickness=1)

    return frame


def distance(x1, y1, x2, y2):
    dist = math.sqrt((x2-x1)**2 + (y2-y1)**2)
    return dist

def rect_in_circle(cX, cY, objects):
    for item in objects:
        if distance(cX, cY, item['x'], item['y']) < item['r']:
            return True
    return False

def top_frame_process(frame, camParams, cvParams, cvSettings):
    objects = {}
    # Convert to grayscale
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # Correct Contrast
    gray=cv2.convertScaleAbs(gray, alpha=cvSettings['contrast'], beta=cvSettings['brightness'])

    # Blur
    #gray_blurred = cv2.blur(gray, (1, 1))
    gray_blurred = cv2.blur(gray, cvSettings['blur'])

    #thresh = cv2.adaptiveThreshold(gray_blurred,255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY,11,2)
    thresh = cv2.adaptiveThreshold(gray_blurred,255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, cvSettings['adaptiveThreshold_blockSize'], cvSettings['adaptiveThreshold_C'])

    if cvParams['img_type'] == 0:
        frameDisplay = frame
    else:
        frameDisplay = thresh

    ################
    # Detect circles
    ################
    objects['circle'] = []
    detected_circles = cv2.HoughCircles(gray_blurred, 
                    cv2.HOUGH_GRADIENT,
                    dp=1.0,
                    minDist = cvParams['circle']['dist-min'],  # Минимальное расстояние между центрами кругов
                    param1=300,  # Порог для градиентного детектора
                    param2=20,  # Порог для акцепта круга
                    minRadius = cvParams['circle']['r-min'],  # Минимальный радиус круга
                    maxRadius = cvParams['circle']['r-max']
        )

    # Draw circles that are detected.
    if detected_circles is not None:

        # Convert the circle parameters a, b and r to integers.
        detected_circles = np.uint16(np.around(detected_circles))

        for pt in detected_circles[0, :]:
            a, b, r = pt[0], pt[1], pt[2]

            # Draw the circumference of the circle.
            frameDisplay = cv2.circle(frameDisplay, (a, b), r, (0, 255, 0), 2)

            # Draw a small circle (of radius 1) to show the center.
            frameDisplay = cv2.circle(frameDisplay, (a, b), 1, (0, 0, 255), 2)

            #print(a, b)
            objects['circle'].append({'x': a.astype(float), 'y': b.astype(float), 'r': r.astype(float)})
        
    ################


    ################
    # Detect rectangle
    ################
    objects['rect'] = []
    #contours, hierarchy = cv2.findContours(thresh, 1, 2)
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

    if contours is not None:
        for cnt in contours:
            #x1,y1 = cnt[0][0]
            approx = cv2.approxPolyDP(cnt, 0.1*cv2.arcLength(cnt, True), True)
            if len(approx) == 4:
                x, y, w, h = cv2.boundingRect(approx)

                #if (w > 20) and (w < 50):
                if (w > cvParams['rectangle']['w-min']) and (w < cvParams['rectangle']['w-max']) and (h > cvParams['rectangle']['h-min']) and (h < cvParams['rectangle']['h-max']):
                    rect = cv2.minAreaRect(approx)
                    box = cv2.boxPoints(rect)

                    box = np.intp(box)
                    #frameDisplay = cv2.drawContours(frameDisplay, [box], 0, (0,0,255), 2)

                    M = cv2.moments(cnt)
                    cX = int(M["m10"] / M["m00"])
                    cY = int(M["m01"] / M["m00"])
                    #cv2.circle(frameDisplay, (cX, cY), 1, (0, 255, 0), 2)

                    rows,cols = frame.shape[:2]
                    [vx,vy,x,y] = cv2.fitLine(cnt, cv2.DIST_L2,0,0.01,0.01)
                    #lefty = int((-x*vy/vx) + y)
                    #righty = int(((cols-x)*vy/vx)+y)
                    #cv2.line(frameDisplay,(cols-1,righty),(0,lefty),(0,255,0),2)

                    dx = cols/2 - cX
                    dy = rows/2 - cY
                    #text1 = "dx: {dx:.2f}, dy: {dy:.2f}".format(dx=dx, dy=dy)
                    #frameDisplay = cv2.putText(frameDisplay, text1, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1)


                    x_axis      = np.array([1, 0])    # unit vector in the same direction as the x axis
                    your_line   = np.array([vx, vy])  # unit vector in the same direction as your line
                    dot_product = np.dot(x_axis, your_line)
                    angle_2_x   = np.arccos(dot_product)
                    angle = math.degrees(angle_2_x[0])

                    #print (angle)
                    #text2 = "angle: {angle:.2f}".format(angle = 2.18)
                    #frameDisplay = cv2.putText(frameDisplay, text2, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1)
                    if rect_in_circle(cX, cY, objects['circle']) == False :
                        frameDisplay = cv2.drawContours(frameDisplay, [box], 0, (0,0,255), 2)
                        frameDisplay = cv2.circle(frameDisplay, (cX, cY), 1, (0, 255, 0), 2)
                        objects['rect'].append({'cx': cX, 'cy': cY, 'angle': angle, 'box': box.tolist()})
    ################


    ################
    # Detect Text
    ################
    # Simple image to string
    if cvParams['text'] == True:
        objects['text'] = []
        data = pytesseract.image_to_data(frame, output_type='dict')
        boxes = len(data['level'])
        for i in range(boxes ):
            (x, y, w, h, text) = (data['left'][i], data['top'][i], data['width'][i], data['height'][i], data['text'][i])
            if (text != ""):
                print (x, y, w, h, text)
                frameDisplay = cv2.rectangle(frameDisplay, (x, y), (x+w, y+h), (200, 255, 200), 3)
                objects['text'].append({'x': x, 'y': y, 'w': w, 'h':h, 'text': text})

    return frameDisplay, objects

def bottom_frame_process(frame_src, camParams, cvParams, cvSettings, negative=False):
    objects = {}

    frame = frame_src.copy()

    frame[:, :, 2] = 255
    frame[:, :, 0] = 255

    if negative == True:
        frame = cv2.bitwise_not(frame)

    # Correct Contrast
    frame = cv2.convertScaleAbs(frame, alpha=cvSettings['contrast'], beta=cvSettings['brightness'])

    # Blur
    #gray_blurred = cv2.blur(gray, (1, 1))
    blurred = cv2.blur(frame, cvSettings['blur'])

    # Convert to grayscale
    gray_blurred = cv2.cvtColor(blurred, cv2.COLOR_BGR2GRAY)

    thresh = cv2.adaptiveThreshold(gray_blurred,255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, cvSettings['adaptiveThreshold_blockSize'], cvSettings['adaptiveThreshold_C'])

    if cvParams['img_type'] == 0:
        frameDisplay = frame_src
    else:
        frameDisplay = thresh

    #contours, hierarchy = cv2.findContours(thresh, 1, 2)
    # RETR_EXTERNAL, RETR_LIST, RETR_CCOMP, RETR_TREE
    # CHAIN_APPROX_NONE, CHAIN_APPROX_SIMPLE
    contours, hierarchy = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)

    rows, cols = frame.shape[:2]
    if contours is not None:
        min_d = 99999
        res_contour = None
        res_aprox = None
        screenCx = cols/2
        screenCy = rows/2
        for cnt in contours:
            approx = cv2.approxPolyDP(cnt, 0.0*cv2.arcLength(cnt, True), True)
            x, y, w, h = cv2.boundingRect(approx)

            if w > cvParams['contour']['w-min'] and w < cvParams['contour']['w-max'] and h > cvParams['contour']['h-min'] and h < cvParams['contour']['h-max']:
                M = cv2.moments(cnt)
                cX = int(M["m10"] / M["m00"])
                cY = int(M["m01"] / M["m00"])
                dX = screenCx - cX
                dY = screenCy - cY
                dist = math.sqrt(dX**2 + dY**2)

                # Find the contour closest to the center
                if dist < w/2 and dist < h/2:
                    if dist < min_d:
                        min_d = dist
                        res_contour = cnt
                        res_aprox = approx

        if res_contour is not None:
            rect = cv2.minAreaRect(res_aprox)
            box = cv2.boxPoints(rect)

            # Draw contour
            box = np.intp(box)
            frameDisplay = cv2.drawContours(frameDisplay, [box], 0, (0,0,255), 2)

            # Draw Center
            M = cv2.moments(res_contour)
            cX = int(M["m10"] / M["m00"])
            cY = int(M["m01"] / M["m00"])
            frameDisplay = cv2.circle(frameDisplay, (cX, cY), 1, (0, 255, 0), 2)

            # Angle to rotate
            sorted_points = sorted(box, key=lambda x: x[1])[:2] # get two top points
            sorted_points = sorted(sorted_points, key=lambda x: x[0]) # sort points from left to right
            
            asin_top_line = (sorted_points[1][1] - sorted_points[0][1]) / (sorted_points[1][0] - sorted_points[0][0])
            angle_degrees = math.degrees(math.atan(asin_top_line))

            # Draw fitLine
            lefty = int(cY - cX*asin_top_line)
            righty = int(cY + (cols-cX)*asin_top_line)
            frameDisplay = cv2.line(frameDisplay,(cols-1,righty),(0,lefty),(0,255,0), 1)

            dx = (cols/2) - cX
            dy = (rows/2) - cY

            objects['object'] = {'dx': dx, 'dy': dy, 'angle_degrees': angle_degrees}
        else:
            objects['object'] = {'dx': None, 'dy': None, 'angle_degrees': None}


        ################
        # Detect Nozzle (circle)
        ################
        objects['nozzle'] = {'dx': None, 'dy': None}
        detected_circles = cv2.HoughCircles(gray_blurred, 
                        cv2.HOUGH_GRADIENT,
                        dp=1.0,
                        minDist = cvParams['circle']['dist-min'],  # Минимальное расстояние между центрами кругов
                        param1=300,  # Порог для градиентного детектора
                        param2=20,  # Порог для акцепта круга
                        minRadius = cvParams['circle']['r-min'],  # Минимальный радиус круга
                        maxRadius = cvParams['circle']['r-max']
            )

        # Draw circles that are detected.
        if detected_circles is not None:
            # Convert the circle parameters a, b and r to integers.
            detected_circles = np.uint16(np.around(detected_circles))

            # Find circle with minimal radius
            min_r = 9999
            min_circle = None
            for pt in detected_circles[0, :]:
                x, y, r = pt[0], pt[1], pt[2]
                if r < min_r:
                    min_r = r
                    min_circle = pt

            if min_circle is not None:
                x, y, r = min_circle[0], min_circle[1], min_circle[2]
                # Draw the circumference of the circle.
                frameDisplay = cv2.circle(frameDisplay, (x, y), r, (0, 255, 0), 2)

                # Draw a small circle (of radius 1) to show the center.
                frameDisplay = cv2.circle(frameDisplay, (x, y), 1, (0, 0, 255), 2)

                dx = (cols/2) - x
                dy = (rows/2) - y
                objects['nozzle'] = {'dx': dx, 'dy': dy}
        ################
    return frameDisplay, objects


def frame_process(num, frame, camParams, cvParams, cvSettings):
    if num == 0:
        return top_frame_process(frame, camParams, cvParams, cvSettings)
    if num == 1: 
        #return bottom_frame_process(frame, camParams, cvParams, cvSettings)
        frameDisplay, objects = bottom_frame_process(frame, camParams, cvParams, cvSettings)
        
        # If no found any objects and nozzle then try found it on the inverted (negative) image. May be object is white.
        if objects['object']['dx'] == None and objects['nozzle']['dx'] == None:
            frameDisplay, objects = bottom_frame_process(frame, camParams, cvParams, cvSettings, True)
        
        return frameDisplay, objects