extract_shapes.py

#!/usr/bin/env python
"""Find shapes in a card image."""

import os
import sys
import cv2
import numpy as np
from common import (
    SHAPES_OUT_DIR,
    display_im,
    write_im,
    clean_make_dir,
    rectify,
    scale_points,
)

# min channel cutoff for the threshold filter
THRESH_MIN = 180

OUT_WIDTH = 100
OUT_HEIGHT = 200


def extract_shapes_from_im(
    card_im, out_w=OUT_WIDTH, out_h=OUT_HEIGHT, display_shapes=False
):
    orig_im = card_im
    im = cv2.cvtColor(orig_im, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(im, (1, 1), 1000)
    flag, thresh = cv2.threshold(blur, THRESH_MIN, 255, cv2.THRESH_BINARY)

    # invert, otherwise RETR_EXTERNAL makes the whole card the largest contour
    thresh = cv2.bitwise_not(thresh)

    # `image` is the thrown away value
    _, contours, hierarchy = cv2.findContours(
        thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    # sort contours by largest volume
    contours = sorted(contours, key=cv2.contourArea, reverse=True)

    shapes = []
    for i in range(min(3, len(contours))):
        rectangle = cv2.minAreaRect(contours[i])
        corners = cv2.boxPoints(rectangle)
        corners = rectify(corners, portrait=True)
        corners = scale_points(corners, 1.15)

        # writes dots for corners of bounding boxes on original image
        # for point in corner:
        # cv2.circle(orig_im, (point[0], point[1]), 0, (0,0,255), im.shape[0]/50)

        # create an image of just the shape
        h = np.array(
            [[0, 0], [out_w - 1, 0], [out_w - 1, out_h - 1], [0, out_h - 1]], np.float32
        )
        transform = cv2.getPerspectiveTransform(corners, h)
        warp = cv2.warpPerspective(orig_im, transform, (out_w, out_h))
        if display_shapes:
            display_im(warp)
        shapes.append(warp)

    return shapes


def extract_shapes_from_file(
    card_file, out_w=OUT_WIDTH, out_h=OUT_HEIGHT, display_shapes=False
):
    """Given a card image file, cut out and return the 1, 2, or 3 shapes
  on the card. Returns a list of lists of 4 points, corner coordinates of
  each bounding box.
  """
    # a lot of this function is the same as card_finder.py's find_cards,
    # TODO refactor for both to inherit from a common function
    orig_im = cv2.imread(card_file, 1)
    return extract_shapes_from_im(
        orig_im, out_w=out_w, out_h=out_h, display_shapes=display_shapes
    )


def write_shape_with_label(shape, input_card_path, out_dir=SHAPES_OUT_DIR):
    """Writes shape to a filename built from a labeled card file, replacing
  its number with 'single'.
  """
    if not os.path.exists(out_dir):
        os.mkdir(out_dir)
    tokens = input_card_path.split("/")[-1].split("-")
    tokens[1] = "single"
    new_filename = "-".join(tokens)
    print(new_filename)
    write_im(shape, new_filename, out_dir=out_dir)


def main():
    card_files = sys.argv[1:]
    for card_file in card_files:
        shapes = extract_shapes_from_file(card_file, display_shapes=False)
        # only works with labeled files
        write_shape_with_label(shapes[0], card_file)


if __name__ == "__main__":
    main()