HTequi/blob-simulation/detection/detect.py

import argparse
import json
from detection.utils import *
from os.path import splitext, basename, join


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("-i", "--input", required=True, help="input image")
    ap.add_argument("-s", "--scale", type=float, default=0.10, help="scale images by this factor (default: x0.1)")
    ap.add_argument("-c", "--config", type=str, default="config.json",
                    help="name file to load config (default: config.json)")
    ap.add_argument("-o", "--output", type=str, help="give a directory name to save the game files in it")
    ap.add_argument("--hide", action='store_true', default=False, help="hide images")

    args = ap.parse_args()
    with open(args.config, 'r') as file:
        config = json.load(file)

    orig, blob_mask, blob, food_mask, food_img = detect(args.input, config)
    discrete_img, discrete_blob, discrete_food_list = discretize(blob, food_mask, config['Discrete Width'],
                                                                             config['Discrete Height'])

    if args.output is not None:
        filename = splitext(basename(args.input))[0]
        dir = args.output
        save(join(dir, filename), discrete_img, discrete_blob, discrete_food_list)

    if not args.hide:
        print_results(orig, blob_mask, blob, food_mask, food_img, discrete_img, args.scale)


def detect(input_file, config):

    img = cv2.imread(input_file)
    height, width, _ = img.shape

    aspect_ratio = config["Aspect Ratio"]

    height = int(width*aspect_ratio)

    """ Resize image to limits in config file """
    limits = np.array(config['Limits'])
    transform_mat = cv2.getPerspectiveTransform(np.float32(limits), np.float32(
        [[0, 0], [width, 0], [width, height], [0, height]]))
    img = cv2.warpPerspective(img, transform_mat, (width, height))

    """ Prepare upper and lower mask board """
    upper_mask = np.zeros(img.shape[0:2], np.uint8)
    lower_mask = np.zeros(img.shape[0:2], np.uint8)

    upper_mask = cv2.rectangle(upper_mask, (0, 0), (width, int(height/2)), 255, thickness=cv2.FILLED)
    lower_mask = cv2.rectangle(lower_mask, (0, int(height / 2)), (width, height), 255, thickness=cv2.FILLED)

    """ Find blob """
    sat = saturation(img)
    # sum = mean_image([satA, satB])
    blob_mask = find_blob(sat)
    blob = cv2.bitwise_and(img, img, mask=blob_mask)

    """ Print blob information """
    print("Blob covering:")
    print("\t{:.2f}% of the board.".format(mean_percent_value(blob_mask)))
    print("\t{:.2f}% of the upper board.".format(
        mean_percent_value(cv2.bitwise_and(blob_mask, blob_mask, mask=upper_mask), img_ratio=0.5)))
    print("\t{:.2f}% of the lower board.".format(
        mean_percent_value(cv2.bitwise_and(blob_mask, blob_mask, mask=lower_mask), img_ratio=0.5)))

    """ Find food """
    food_list, food_mask, food_img = find_food(img, config['Min Food Size'], config['Low Food Color'], config['High Food Color'])

    """ Print food information """
    print("Total food discovered: " + str(len(food_list)))
    # for i, food in enumerate(food_list):
        # print("\tFood N°" + str(i) + ": " + str(food))

    return img, blob_mask, blob, food_mask, food_img


def print_results(orig, blob_mask, blob, food_mask, food, discrete, scale=1.0):
    padding = 35
    nbr_width = 2
    nbr_height = 3
    font = cv2.FONT_HERSHEY_SIMPLEX
    fontsize = 0.45
    thickness = 1

    scaled_height = int(orig.shape[0]*scale)
    scaled_width = int(orig.shape[1]*scale)

    pad = np.zeros((scaled_height, padding, orig.shape[2]), dtype=np.uint8)
    line_pad = np.zeros((padding, (scaled_width + padding) * nbr_width + padding, orig.shape[2]), dtype=np.uint8)
    print_img = cv2.resize(orig, (scaled_width, scaled_height))

    middle = ((0, int(scaled_height/2)), (scaled_width, int(scaled_height/2)))
    cv2.line(print_img, middle[0], middle[1], (0, 255, 0), thickness=1)
    cv2.putText(print_img, 'Mid Line', (middle[0][0] + 5, middle[0][1] - 5),
                font, fontsize, (0, 255, 0), thickness, cv2.LINE_AA)

    print_blob_mask = cv2.resize(cv2.cvtColor(blob_mask, cv2.COLOR_GRAY2BGR), (scaled_width, scaled_height))
    print_blob = cv2.resize(blob, (scaled_width, scaled_height))
    print_food_mask = cv2.resize(cv2.cvtColor(food_mask, cv2.COLOR_GRAY2BGR), (scaled_width, scaled_height))
    print_food = cv2.resize(food, (scaled_width, scaled_height))
    print_discrete = cv2.resize(discrete, (scaled_width, scaled_height))

    concat_line1 = np.concatenate((pad, print_img, pad, print_discrete, pad), axis=1)
    concat_line2 = np.concatenate((pad, print_blob_mask, pad, print_blob, pad), axis=1)
    concat_line3 = np.concatenate((pad, print_food_mask, pad, print_food, pad), axis=1)

    aggregate = np.concatenate((line_pad, concat_line1, line_pad, concat_line2, line_pad, concat_line3, line_pad))

    cv2.putText(aggregate, 'Original:',
                (0 * (scaled_width + padding) + padding + 5, 0 * (scaled_height + padding) + padding - 5),
                font, fontsize, (255, 255, 255), thickness, cv2.LINE_AA)
    cv2.putText(aggregate, 'Discrete:',
                (1 * (scaled_width + padding) + padding + 5, 0 * (scaled_height + padding) + padding - 5),
                font, fontsize, (255, 255, 255), thickness, cv2.LINE_AA)
    cv2.putText(aggregate, 'Blob Mask:',
                (0 * (scaled_width + padding) + padding + 5, 1 * (scaled_height + padding) + padding - 5),
                font, fontsize, (255, 255, 255), thickness, cv2.LINE_AA)
    cv2.putText(aggregate, 'Blob:',
                (1 * (scaled_width + padding) + padding + 5, 1 * (scaled_height + padding) + padding - 5),
                font, fontsize, (255, 255, 255), thickness, cv2.LINE_AA)
    cv2.putText(aggregate, 'Food Mask:',
                (0 * (scaled_width + padding) + padding + 5, 2 * (scaled_height + padding) + padding - 5),
                font, fontsize, (255, 255, 255), thickness, cv2.LINE_AA)
    cv2.putText(aggregate, 'Food Regions:',
                (1 * (scaled_width + padding) + padding + 5, 2 * (scaled_height + padding) + padding - 5),
                font, fontsize, (255, 255, 255), thickness, cv2.LINE_AA)

    cv2.imshow("Results", aggregate)
    print("\nPress any key...")
    cv2.waitKey(0)


def discretize(blob_img, food_mask, width, height):
    img_height, img_width, _ = blob_img.shape

    discrete_blob = cv2.resize(blob_img, (width, height), interpolation=cv2.INTER_NEAREST)

    discrete_food = cv2.resize(food_mask, (width, height), interpolation=cv2.INTER_NEAREST)

    discrete_food_list = []
    for x in range(height):
        for y in range(width):
            if discrete_food[x, y] != 0:
                discrete_food_list.append((y, x))

    height, width, _ = discrete_blob.shape
    discrete_blob = cv2.cvtColor(discrete_blob, cv2.COLOR_BGR2GRAY)

    # If discrete blob has to be connected, used this :
    # contours, hierarchy = cv2.findContours(discrete_blob, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    # c = max(contours, key=cv2.contourArea)
    # mask = np.zeros(discrete_blob.shape, np.uint8)
    # cv2.drawContours(mask, [c], -1, 255, cv2.FILLED)
    # discrete_blob = cv2.bitwise_and(discrete_blob, discrete_blob, mask=mask)

    discrete_blob_bgr = cv2.cvtColor(discrete_blob, cv2.COLOR_GRAY2BGR)
    discrete_img = cv2.resize(discrete_blob_bgr, (0, 0), fx=10, fy=10, interpolation=cv2.INTER_NEAREST)
    for (x, y) in discrete_food_list:
        cv2.rectangle(discrete_img, (x * 10, y * 10), ((x + 1) * 10, (y + 1) * 10), (0, 255, 0), thickness=cv2.FILLED)
        if discrete_blob[y, x] != 0:
            cv2.drawMarker(discrete_img, (x * 10 + 5, y * 10 + 5), (255, 255, 255), thickness=2, markerSize=9,
                           markerType=cv2.MARKER_TILTED_CROSS)

    return discrete_img, discrete_blob, discrete_food_list


def save(filename, discrete_img, discrete_blob, discrete_food_list):

    height, width = discrete_blob.shape

    board_str = str(width) + ' ' + str(height) + '\n'

    discrete_food = np.zeros(discrete_blob.shape, dtype=bool)
    for (x, y) in discrete_food_list:
        discrete_food[y, x] = True

    for x in range(height):
        for y in range(width):
            board_str += "{:d},{:d},{} ".format(discrete_blob[x, y] != 0, discrete_food[x,y], discrete_blob[x, y])
        board_str = board_str[:-1]
        board_str += "\n"
    board_str = board_str[:-1]

    with open(filename + ".board", 'w') as board_file:
        board_file.write(board_str)

    cv2.imwrite(filename + ".jpg", discrete_img)


if __name__ == "__main__":
    main()