Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- """
- file: extract.py
- auth: Peb Ruswono Aryan
- date: 02.10.2019
- desc: extract shapes (rectangle,triangle,circle) and raw attributes (color, location and size)
- for used in kandinsky test
- output tab separated values in the stdout
- """
- import numpy as np
- import cv2
- import argparse
- def str2color(s):
- if s=='red':
- return (0,0,255)
- elif s=='blue':
- return (255,0,0)
- elif s=='yellow':
- return (0,255,255)
- else:
- return (255,255,255)
- def extract_shape(imgray, color, display=False):
- ret, thresh = cv2.threshold(imgray, 127, 255, 0)
- contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
- if display:
- shapes = np.zeros([imgray.shape[0], imgray.shape[1], 3], dtype=np.uint8)
- objects = []
- for obj in contours:
- if len(obj)>15:
- #circle
- shape = 'circle'
- elif len(obj)<13:
- #rectangle
- shape = 'rect'
- else:
- shape = 'tri'
- # triangle
- loc = ((obj.max(axis=0)+obj.min(axis=0))//2)[0]
- size = (obj.max(axis=0)-obj.min(axis=0))[0]
- objects.append([color, shape,loc[0],loc[1],size[0],size[1]])
- if display:
- cv2.drawContours(shapes, [obj], 0, str2color(color), 1)
- if display:
- cv2.imshow("shape_"+color, shapes)
- return objects
- def detect_red(img):
- red = [0, 0, 255]
- mask = img == np.array(red)
- rimg = img.copy()
- rimg[~mask] = 0
- gimg = cv2.cvtColor(rimg, cv2.COLOR_BGR2GRAY)
- gimg[gimg>0] = 255
- return gimg
- def detect_blue(img):
- blue = [255, 0, 0]
- mask = img == np.array(blue)
- rimg = img.copy()
- rimg[~mask] = 0
- gimg = cv2.cvtColor(rimg, cv2.COLOR_BGR2GRAY)
- gimg[gimg>0] = 255
- return gimg
- def detect_yellow(img):
- red = [0, 0, 255]
- green = [0, 255, 0]
- mask = (img == np.array(green))
- rimg = img.copy()
- rimg[~mask] = 0
- gimg = cv2.cvtColor(rimg, cv2.COLOR_BGR2GRAY)
- gimg[gimg>0] = 255
- return gimg
- if __name__=="__main__":
- parser = argparse.ArgumentParser()
- parser.add_argument('-display', action='store_true', default=False, help='display images in separate window')
- parser.add_argument('inputfile')
- args = parser.parse_args()
- img = cv2.imread(args.inputfile)
- redimg = detect_red(img)
- bluimg = detect_blue(img)
- yelimg = detect_yellow(img)
- if args.display:
- cv2.imshow("", img)
- cv2.imshow("red", redimg)
- cv2.imshow("blue", bluimg)
- cv2.imshow("yellow", yelimg)
- allobj = []
- allobj += extract_shape(redimg, 'red', args.display)
- allobj += extract_shape(bluimg, 'blue', args.display)
- allobj += extract_shape(yelimg, 'yellow', args.display)
- if args.display:
- cv2.waitKey(0)
- print('\t'.join(['color', 'shape', 'x', 'y', 'w', 'h']))
- for obj in allobj:
- print('\t'.join([str(x) for x in obj]))
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement