extract_raw_kandinsky.py

"""
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]))