Velaxers

AUV getSmallGateGrid(image)

Oct 16th, 2020
865
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. import cv2
  2. import numpy as np
  3. import matplotlib as plt
  4. from matplotlib import pyplot as plt
  5. import PIL
  6. from PIL import Image
  7. from PIL import ImageEnhance
  8. from pylab import rcParams
  9.  
  10. def getGateGridNumber(image):
  11.     imageEnhanced = PIL.ImageEnhance.Brightness(image).enhance(0.9)
  12.     imageEnhanced = PIL.ImageEnhance.Contrast(imageEnhanced).enhance(2)
  13.     imageEnhanced = PIL.ImageEnhance.Sharpness(imageEnhanced).enhance(3.2)
  14.  
  15.     imageCV = cv2.cvtColor(np.array(imageEnhanced), cv2.COLOR_RGB2BGR)
  16.  
  17.     b, g, r = cv2.split(imageCV)
  18.     b = cv2.medianBlur(b,5)
  19.  
  20.     kernel = np.ones((5,5), np.uint8)
  21.     b = cv2.erode(b, kernel, iterations=5)
  22.     b = cv2.dilate(b, kernel, iterations=5)
  23.  
  24.     rcParams['figure.figsize'] = 10, 12
  25.     edges = cv2.Canny(b,
  26.                     threshold1=100, ## try different values here
  27.                     threshold2=100) ## try different values here
  28.                    
  29.     #plt.title('Edge Detection')
  30.     #plt.imshow(edges)
  31.     #plt.show()
  32.  
  33.     contours, hierarchy = cv2.findContours(edges,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
  34.     try: hierarchy = hierarchy[0]
  35.     except: hierarchy = []
  36.  
  37.     imageOutput = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
  38.  
  39.     height, width = edges.shape
  40.     min_x, min_y = width, height
  41.     max_x = max_y = 0
  42.  
  43.     contoursFiltered = []
  44.     # computes the bounding box for the contour, and draws it on the frame,
  45.     i = 0
  46.     rectanglesWidths = []
  47.     rectanglesHeights = []
  48.     for contour, hier in zip(contours, hierarchy):
  49.         (x,y,w,h) = cv2.boundingRect(contour)
  50.         min_x, max_x = min(x, min_x), max(x+w, max_x)
  51.         min_y, max_y = min(y, min_y), max(y+h, max_y)
  52.         color = (0, 255, 0) if (i == 0) else (0,0,255)
  53.         if w > 100 and h > 50:
  54.             h = 250
  55.             cv2.rectangle(imageOutput, (x,y), (x+w,y+h), color, 2)
  56.             contoursFiltered.append(contour)
  57.             i += 1
  58.             # Saving widths and heights to later on put text on the center of the rectangles
  59.             width = x+(w/2)
  60.             height = y+(h/2)
  61.             rectanglesWidths.append(width)
  62.             rectanglesHeights.append(height)
  63.  
  64.     if max_x - min_x > 0 and max_y - min_y > 0:
  65.         cv2.rectangle(imageOutput, (min_x +50, min_y +150), (max_x, max_y), (255, 0, 0), 2)
  66.  
  67.     if len(rectanglesWidths)>=2:
  68.         for i in range(len(rectanglesWidths)):
  69.             if (rectanglesWidths[0] * rectanglesHeights[0] > rectanglesWidths[1] * rectanglesHeights[1]):
  70.                 cv2.putText(imageOutput, text= 'Big Gate', org=(int(rectanglesWidths[0]-90),int(rectanglesHeights[0])),
  71.                     fontFace= cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,0),
  72.                     thickness=2, lineType=cv2.LINE_AA)
  73.                 cv2.putText(imageOutput, text= 'Small Gate', org=(int(rectanglesWidths[1]-90),int(rectanglesHeights[1])),
  74.                     fontFace= cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,0),
  75.                     thickness=2, lineType=cv2.LINE_AA)
  76.             else:
  77.                 cv2.putText(imageOutput, text= 'Big Gate', org=(int(rectanglesWidths[1]-90),int(rectanglesHeights[1])),
  78.                     fontFace= cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,0),
  79.                     thickness=2, lineType=cv2.LINE_AA)
  80.                 cv2.putText(imageOutput, text= 'Small Gate', org=(int(rectanglesWidths[0]-90),int(rectanglesHeights[0])),
  81.                     fontFace= cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=(0,0,0),
  82.                     thickness=2, lineType=cv2.LINE_AA)
  83.  
  84.     #plt.title("Gate Detection")
  85.     #plt.imshow(cv2.cvtColor(imageOutput, cv2.COLOR_BGR2RGB))
  86.     #plt.show()
  87.  
  88.     if len(rectanglesWidths)>=2:
  89.         for i in range(len(rectanglesWidths)):
  90.             if (rectanglesWidths[0] * rectanglesHeights[0] < rectanglesWidths[1] * rectanglesHeights[1]):
  91.                 centerX = rectanglesWidths[0]
  92.                 centerY = rectanglesHeights[0]
  93.             else:
  94.                 centerX = rectanglesWidths[1]
  95.                 centerY = rectanglesHeights[1]
  96.  
  97.     #print(centerX, centerY)
  98.  
  99.     #Code to detect which grid (1 to 9)
  100.     #Send the integer to the C++ node
  101.     height, width = edges.shape
  102.     #print(height, width)
  103.     hThird = height / 3
  104.     wThird = width / 3
  105.     if (centerX >= 0 and centerY >= 0 and centerX <= wThird and centerY <=hThird):
  106.         grid = "1"
  107.     elif (centerX >= wThird and centerY >= 0 and centerX <= wThird*2 and centerY <=hThird):
  108.         grid = "2"
  109.     elif (centerX >= wThird*2 and centerY >= 0 and centerX <= wThird*3 and centerY <=hThird):
  110.         grid = "3"
  111.     elif (centerX >= 0 and centerY >= hThird and centerX <= wThird and centerY <=hThird*2):
  112.         grid = "4"
  113.     elif (centerX >= wThird and centerY >= hThird and centerX <= wThird*2 and centerY <=hThird*2):
  114.         grid = "5"
  115.     elif (centerX >= wThird*2 and centerY >= hThird and centerX <= wThird*3 and centerY <=hThird*2):
  116.         grid = "6"
  117.     elif (centerX >= 0 and centerY >= hThird*2 and centerX <= wThird*1 and centerY <=hThird*3):
  118.         grid = "7"
  119.     elif (centerX >= wThird and centerY >= hThird*2 and centerX <= wThird*2 and centerY <=hThird*3):
  120.         grid = "8"
  121.     elif (centerX >= wThird*2 and centerY >= hThird*2 and centerX <= wThird*3 and centerY <=hThird*3):
  122.         grid = "9"
  123.  
  124.     return grid
  125.  
  126.  
  127. #### Main starts here
  128. #image = Image.open("C:/Users/Lenovo/Desktop/Vortex AUV/testing images/noise1.png")
  129. imagePath = Image.open("C:/Users/Lenovo/Desktop/Vortex AUV/Gate video frames/Edge Data/createData8/35.png")
  130. print(getGateGridNumber(imagePath))
RAW Paste Data