Advertisement
dan-masek

Untitled

Jan 16th, 2020
381
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
Python 5.95 KB | None | 0 0
  1. from cv2_41 import cv2
  2. import numpy as np
  3.  
  4. # rotate_bound: helper function that rotates the image adds some padding to avoid cutting off parts of it
  5. # reference: https://www.pyimagesearch.com/2017/01/02/rotate-images-correctly-with-opencv-and-python/
  6. def rotate_bound(image, angle):
  7.     # grab the dimensions of the image and then determine the center
  8.     (h, w) = image.shape[:2]
  9.     (cX, cY) = (w // 2, h // 2)
  10.  
  11.     # grab the rotation matrix (applying the negative of the angle to rotate clockwise), then grab the sine and cosine
  12.     # (i.e., the rotation components of the matrix)
  13.     M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
  14.     cos = np.abs(M[0, 0])
  15.     sin = np.abs(M[0, 1])
  16.  
  17.     # compute the new bounding dimensions of the image
  18.     nW = int(np.multiply(h, sin) + np.multiply(w, cos))
  19.     nH = int(np.multiply(h, cos) + np.multiply(w, sin))
  20.  
  21.     # adjust the rotation matrix to take into account translation
  22.     M[0, 2] += (nW / 2) - cX
  23.     M[1, 2] += (nH / 2) - cY
  24.  
  25.     # perform rotation and return the image (white background) along with the Rotation Matrix
  26.     return cv2.warpAffine(image, M, (nW, nH), borderValue=(255,255,255)), M
  27.  
  28.  
  29. # Step 1 - Load images
  30. input_img = cv2.imread("target.png", cv2.IMREAD_GRAYSCALE)
  31. template_img = cv2.imread("template.png", cv2.IMREAD_GRAYSCALE)
  32. matches_dbg_img = cv2.cvtColor(input_img, cv2.COLOR_GRAY2BGR) # for debugging purposes
  33.  
  34. # Step 2 - Generate some ROIs
  35. # each ROI contains the x,y,w,h and angle (degree) to rotate the box and make its M appear horizontal
  36. roi_w = 26
  37. roi_h = 26
  38.  
  39. roi_list = []
  40. roi_list.append((112, 7, roi_w, roi_h, 0))
  41. roi_list.append((192, 36, roi_w, roi_h, -45))
  42. roi_list.append((227, 104, roi_w, roi_h, -90))
  43. roi_list.append((195, 183, roi_w, roi_h, -135))
  44. roi_list.append((118, 216, roi_w, roi_h, -180))
  45. roi_list.append((49, 196, roi_w, roi_h, -225))
  46. roi_list.append((10, 114, roi_w, roi_h, -270))
  47. roi_list.append((36, 41, roi_w, roi_h, -315))
  48.  
  49. # debug: draw green ROIs
  50. rois_dbg_img = cv2.cvtColor(input_img, cv2.COLOR_GRAY2BGR)
  51. for roi in roi_list:
  52.     x, y, w, h, angle = roi
  53.     x2 = x + w
  54.     y2 = y + h
  55.     cv2.rectangle(rois_dbg_img, (x, y), (x2, y2), (0,255,0), 2)
  56.  
  57. cv2.imwrite('target_rois.png', rois_dbg_img)
  58. cv2.imshow('ROIs', rois_dbg_img)
  59. cv2.waitKey(0)
  60. cv2.destroyWindow('ROIs')
  61.  
  62.  
  63. # Step 3 - Select a ROI, crop and rotate it, then perform Template Matching
  64. for i, roi in enumerate(roi_list):
  65.     x, y, w, h, angle = roi
  66.     roi_cropped = input_img[y:y+h, x:x+w]
  67.     roi_rotated, M = rotate_bound(roi_cropped, angle)
  68.  
  69.     # debug: display each rotated ROI
  70.     #cv2.imshow('ROIs-cropped-rotated', roi_rotated)
  71.     #cv2.waitKey(0)
  72.  
  73.     # debug: dump roi to the disk (before/after rotation)
  74.     filename = 'target_roi' + str(i)
  75.     cv2.imwrite(filename + '.png', roi_cropped)
  76.     cv2.imwrite(filename + '_rotated.png', roi_rotated)
  77.  
  78.     # perform template matching
  79.     res = cv2.matchTemplate(roi_rotated, template_img, cv2.TM_CCOEFF_NORMED)
  80.     (_, score, _, (pos_x, pos_y)) = cv2.minMaxLoc(res)
  81.     print('TM score=', score)
  82.  
  83.     # Step 4 - When a TM is found, revert the rotation of matched point so that it represents a location in the original image
  84.     # Note: pos_x and pos_y define the location of the matched template in a rotated ROI
  85.     threshold = 0.75
  86.     if (score >= threshold):
  87.  
  88.         # debug in cropped image
  89.         print('find_k_symbol: FOUND pos_x=', pos_x, 'pos_y=', pos_y, 'w=', template_img.shape[1], 'h=', template_img.shape[0])
  90.         rot_output_roi = cv2.cvtColor(roi_rotated, cv2.COLOR_GRAY2BGR)
  91.         cv2.rectangle(rot_output_roi, (pos_x, pos_y), (pos_x + template_img.shape[1], pos_y + template_img.shape[0]), (0, 165, 255), 2) # orange
  92.         cv2.imshow('rot-matched-template', rot_output_roi)
  93.         cv2.waitKey(0)
  94.         cv2.destroyWindow('rot-matched-template')
  95.  
  96.         ###
  97.         # How to convert the location of the matched template (pos_x, pos_y) to points in roi_cropped?
  98.         # (which is the ROI before rotation)
  99.         ###
  100.  
  101.         # extract variables from the rotation matrix
  102.         M_x = M[0][2]
  103.         M_y = M[1][2]
  104.         #print('M_x=', M_x, '\tM_y=', M_y)
  105.         M_cosx = M[0][0]
  106.         M_msinx = M[0][1]
  107.         #print('M_cosx=', M_cosx, '\tM_msinx=', M_msinx)
  108.         M_siny = M[1][0]
  109.         M_cosy = M[1][1]
  110.         #print('M_siny=', M_siny, '\tM_cosy=', M_cosy)
  111.  
  112.         # undo translation:
  113.         dst1_x = pos_x - M_x
  114.         dst1_y = pos_y - M_y
  115.  
  116.         # undo rotation:
  117.         # after this operation, (new_pos_x, new_pos_y) should already be a valid point in the original ROI
  118.         new_pos_x =  M_cosx * dst1_x - M_msinx * dst1_y
  119.         new_pos_y = -M_siny * dst1_x + M_cosy  * dst1_y
  120.  
  121.         # debug: create the bounding rect of the detected symbol in the original input image
  122.         detected_x = x + int(new_pos_x)
  123.         detected_y = y + int(new_pos_y)
  124.         detected_w = template_img.shape[1]
  125.         detected_h = template_img.shape[0]
  126.         detected_rect = (detected_x, detected_y, detected_w, detected_h)
  127.        
  128.         print('find_k_symbol: detected_x=', detected_x, 'detected_y=', detected_y, 'detected_w=', detected_w, 'detected_h=', detected_h)
  129.         print()
  130.  
  131.         bb_points = np.array([
  132.             (detected_x, detected_y)
  133.             , (detected_x + detected_w, detected_y)
  134.             , (detected_x + detected_w, detected_y + detected_h)
  135.             , (detected_x, detected_y + detected_h)
  136.         ])
  137.         Mrot = cv2.getRotationMatrix2D((detected_x, detected_y), angle, 1.0)
  138.        
  139.         rot_points = np.array([np.dot(Mrot, np.array([p[0], p[1], 1]).T) for p in bb_points], np.int32)
  140.         cv2.polylines(matches_dbg_img, [rot_points], True, (0, 165, 255), 2)
  141.  
  142.         #cv2.rectangle(matches_dbg_img, (detected_x, detected_y), (detected_x + detected_w, detected_y + detected_h), (0, 165, 255), 2) # orange
  143.         cv2.imwrite('target_matches.png', matches_dbg_img)
  144.         cv2.imshow('matches', matches_dbg_img)
  145.         cv2.waitKey(0)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement