detect_board_1.py

from __future__ import print_function
from nis import match
import cv2 as cv2
import numpy as np
from matplotlib import pyplot as plt
import argparse
import os

# Parse arguments
parser = argparse.ArgumentParser(description='Code for detecting peeled areas in vegetables')
parser.add_argument('--folder', help='Path to input images.', default='potato')
parser.add_argument('--image', help='Name of image to process. If empty all images are processed.', default='')
args = parser.parse_args()

# Iterate through folder or select specific image
path = os.path.join("./img_data/", args.folder)
images =  next(os.walk(path), (None, None, []))[2] if args.image == '' else [args.image]

# Loop through images and process
for imFile in sorted(images):

    # Load image and convert to gray
    print(f"Processing {imFile }")
    imPath = os.path.join(path, imFile)
    img = cv2.imread(imPath) 
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    if img is None:
        print('Could not open or find the images!')
        exit(0)

    ########## Process image ########## 

    

    # Initiate SIFT detector
    sift = cv2.SIFT_create()

    # Find the keypoints and descriptors with SIFT
    kp1, des1 = sift.detectAndCompute(gray1,None)
    kp2, des2 = sift.detectAndCompute(gray2,None)

    # FLANN parameters
    FLANN_INDEX_KDTREE = 0
    index_params = dict(algorithm = FLANN_INDEX_KDTREE, trees = 5)
    search_params = dict(checks=50)   # or pass empty dictionary

    # Get matches
    flann = cv2.FlannBasedMatcher(index_params,search_params)
    matches = flann.knnMatch(des1,des2,k=2)
    print(f"Matches: {len(matches)}")

    # Need to draw only good matches, so create a mask
    matchesMask = [[0,0] for i in range(len(matches))]
    goodMatches = []

    # Ratio test as per Lowe's paper
    ratioThresh = 1
    for i,(m,n) in enumerate(matches):
        if m.distance < ratioThresh * n.distance:
            matchesMask[i]=[1,0]
            goodMatches.append(m)

    draw_params = dict(matchColor = (0,255,0),
                    singlePointColor = (255,0,0),
                    matchesMask = matchesMask,
                    flags = 0)

    img3 = cv2.drawMatchesKnn(img1,kp1,img2,kp2,matches,None,**draw_params)

    # -------- Draw bounding box ---------
    import math

    def rotate(origin, point, angle):
        """
        Rotate a point counterclockwise by a given angle around a given origin.

        The angle should be given in radians.
        """
        ox, oy = origin
        px, py = point

        qx = ox + math.cos(angle) * (px - ox) - math.sin(angle) * (py - oy)
        qy = oy + math.sin(angle) * (px - ox) + math.cos(angle) * (py - oy)
        return qx, qy

    # Find top two matches
    goodMatches = sorted(goodMatches, key = lambda x:x.distance)
    goodMatches = goodMatches[:2]
    print(f'Good Matches: {goodMatches}')

    # Get angles between top two points in query and train images
    queryPt1 = kp1[goodMatches[0].queryIdx].pt 
    queryPt2 = kp1[goodMatches[1].queryIdx].pt
    trainPt1 = kp2[goodMatches[0].trainIdx].pt
    trainPt2 = kp2[goodMatches[1].trainIdx].pt
    queryAngle =  np.arctan2(queryPt2[0] - queryPt1[0], queryPt2[1] - queryPt1[1])
    print(f'Query Angle: {queryAngle}')
    trainAngle =  np.arctan2(trainPt2[0] - trainPt1[0], trainPt2[1] - trainPt1[1])
    print(f'Train Angle: {trainAngle}')
    AngleDiff =  queryAngle - trainAngle

    # Find center of query image in relation to top two query points
    h, w = img1.shape[:2]
    queryCenter = [w//2, h//2]
    print(f"Height: {h} Width: {w} Center: {queryCenter}")

    queryOffset = [queryCenter[0] - queryPt1[0] , queryCenter[1] - queryPt1[1]]
    print(f'Query Offset: {queryOffset}')

    # Find center in train image using top train point and angle difference
    trainOffset = rotate([0,0], queryOffset, AngleDiff)
    trainCenter = [int(trainPt1[0] + trainOffset[0]), int(trainPt1[1] + trainOffset[1])]
    print(f"Train Center: {trainCenter}")

    # Draw center
    img3 = cv2.circle(img3, [trainCenter[0]  + w,    trainCenter[1]],  10, (255, 0, 0), -1)
    img3 = cv2.circle(img3, [int(trainPt1[0] + w), int(trainPt1[1])],  5, (0, 0, 255), 2)
    img3 = cv2.circle(img3, [int(trainPt2[0] + w), int(trainPt2[1])],  5, (0, 0, 0), 2)
    img3 = cv2.circle(img3, queryCenter,                              10, (255, 0, 0), -1)
    img3 = cv2.circle(img3, [int(queryPt1[0]),     int(queryPt1[1])],  5, (0, 0, 255), 2)
    img3 = cv2.circle(img3, [int(queryPt2[0]),     int(queryPt2[1])],  5, (0, 0, 0), 2)

    plt.imshow(img3,),plt.show()