SHARE
TWEET

Untitled

a guest Oct 21st, 2019 94 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. """
  2. file: extract.py
  3. auth: Peb Ruswono Aryan
  4. date: 02.10.2019
  5. desc: extract shapes (rectangle,triangle,circle) and raw attributes (color, location and size)
  6. for used in kandinsky test
  7. output tab separated values in the stdout
  8. """
  9. import numpy as np
  10. import cv2
  11. import argparse
  12.  
  13. def str2color(s):
  14.     if s=='red':
  15.         return (0,0,255)
  16.     elif s=='blue':
  17.         return (255,0,0)
  18.     elif s=='yellow':
  19.         return (0,255,255)
  20.     else:
  21.         return (255,255,255)
  22.  
  23. def extract_shape(imgray, color, display=False):
  24.     ret, thresh = cv2.threshold(imgray, 127, 255, 0)
  25.     contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
  26.  
  27.     if display:
  28.         shapes = np.zeros([imgray.shape[0], imgray.shape[1], 3], dtype=np.uint8)
  29.  
  30.     objects = []
  31.     for obj in contours:
  32.         if len(obj)>15:
  33.             #circle
  34.             shape = 'circle'
  35.         elif len(obj)<13:
  36.             #rectangle
  37.             shape = 'rect'
  38.         else:
  39.             shape = 'tri'
  40.             # triangle
  41.         loc = ((obj.max(axis=0)+obj.min(axis=0))//2)[0]
  42.         size = (obj.max(axis=0)-obj.min(axis=0))[0]
  43.         objects.append([color, shape,loc[0],loc[1],size[0],size[1]])
  44.         if display:
  45.             cv2.drawContours(shapes, [obj], 0, str2color(color), 1)
  46.     if display:
  47.         cv2.imshow("shape_"+color, shapes)
  48.     return objects
  49.  
  50. def detect_red(img):
  51.     red = [0, 0, 255]
  52.     mask = img == np.array(red)
  53.  
  54.     rimg = img.copy()
  55.     rimg[~mask] = 0
  56.  
  57.     gimg = cv2.cvtColor(rimg, cv2.COLOR_BGR2GRAY)
  58.     gimg[gimg>0] = 255
  59.  
  60.     return gimg
  61.  
  62. def detect_blue(img):
  63.     blue = [255, 0, 0]
  64.     mask = img == np.array(blue)
  65.  
  66.     rimg = img.copy()
  67.     rimg[~mask] = 0
  68.  
  69.     gimg = cv2.cvtColor(rimg, cv2.COLOR_BGR2GRAY)
  70.     gimg[gimg>0] = 255
  71.  
  72.     return gimg
  73.  
  74. def detect_yellow(img):
  75.     red = [0, 0, 255]
  76.     green = [0, 255, 0]
  77.     mask = (img == np.array(green))
  78.  
  79.     rimg = img.copy()
  80.     rimg[~mask] = 0
  81.  
  82.     gimg = cv2.cvtColor(rimg, cv2.COLOR_BGR2GRAY)
  83.     gimg[gimg>0] = 255
  84.  
  85.     return gimg
  86.  
  87. if __name__=="__main__":
  88.     parser = argparse.ArgumentParser()
  89.     parser.add_argument('-display', action='store_true', default=False, help='display images in separate window')
  90.     parser.add_argument('inputfile')
  91.     args = parser.parse_args()
  92.  
  93.     img = cv2.imread(args.inputfile)
  94.  
  95.     redimg = detect_red(img)
  96.     bluimg = detect_blue(img)
  97.     yelimg = detect_yellow(img)
  98.  
  99.     if args.display:
  100.         cv2.imshow("", img)
  101.         cv2.imshow("red", redimg)
  102.         cv2.imshow("blue", bluimg)
  103.         cv2.imshow("yellow", yelimg)
  104.    
  105.     allobj = []
  106.     allobj += extract_shape(redimg, 'red', args.display)
  107.     allobj += extract_shape(bluimg, 'blue', args.display)
  108.     allobj += extract_shape(yelimg, 'yellow', args.display)
  109.  
  110.     if args.display:
  111.         cv2.waitKey(0)
  112.  
  113.     print('\t'.join(['color', 'shape', 'x', 'y', 'w', 'h']))
  114.     for obj in allobj:
  115.         print('\t'.join([str(x) for x in obj]))
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top