romtool/cvtest.py

import cv2
import imutils
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
from matplotlib import colors

RESIZE_RATIO = 0.07
FILENAMES = ["testimg.jpg", "testimg2.jpg", "testimg3.jpg"]

def shrink(img):
    return cv2.resize(img, None, fx=RESIZE_RATIO, fy=RESIZE_RATIO)

def noop(*x):
    pass

def darken_color(img, sat=20, val=225):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    h, s, v = cv2.split(hsv)
    s[s > sat] = 255
    v[v < val] = 0
    hsv2 = cv2.merge((h, s, v))

    return cv2.cvtColor(hsv2, cv2.COLOR_HSV2BGR)

def interactive_edge(filename):
    img = shrink(cv2.imread(filename))

    cv2.createTrackbar("threshold1", "image", 5, 50, noop)
    cv2.createTrackbar("threshold2", "image", 20, 50, noop)

    while handle_key():
        threshold1 = cv2.getTrackbarPos("threshold1", "image")
        threshold2 = cv2.getTrackbarPos("threshold2", "image")
        edged = edge_img(img, threshold1 * 10, threshold2 * 10)
        # contours = cv2.findContours(edged, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        # contours = imutils.grab_contours(contours)
        # shapes = cv2.drawContours(edged.copy(), contours, -1, (255, 255, 255), -1)
        cv2.imshow("image", edged)


def edge_img(orig, t1=50, t2=200):
    img = cv2.GaussianBlur(orig.copy(), (5, 5), 0)
    img = darken_color(img)

    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray_blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    return cv2.Canny(gray_blurred, t1, t2)


def approximate_contour(c):
    peri = cv2.arcLength(c, True)
    return cv2.approxPolyDP(c, 0.02 * peri, True)


def group_contours(contours):
    remaining = contours.copy()
    groups = []
    while len(remaining) > 0:
        group = [remaining[0]]
        groups.append(group)
        grew = False
        remaining = remaining[1:]
        while True:
            nextremaining = []
            l, t, r, b = cv2.boundingRect(remaining[0])
            for c in remaining:
                l2, t2, r2, b2 = cv2.boundingRect(c)
                if r2 < l or l2 > r or t2 > b or b2 < t:
                    # they do not intersect
                    nextremaining.append(c)
                else:
                    # they do intersect
                    grew = True
                    group.append(c)
                    l = min(l, l2)
                    t = min(t, t2)
                    r = max(r, r2)
                    b = max(b, b2)
            remaining = nextremaining
            if not grew:
                break
    return groups

def rects_intersect(rect1, rect2):
    l1, t1, w1, h1 = rect1
    l2, t2, w2, h2 = rect2
    r1 = l1 + w1
    r2 = l2 + w2
    b1 = t1 + h1
    b2 = t2 + h2
    intersect = not (r2 < l1 or l2 > r1 or t2 > b1 or b2 < t1)
    print("1:", rect1, "2:", rect2, intersect)
    return intersect

def merge_contours(contours):
    while True:
        # take the convex hull of all discovered contours, then merge any contours whose bounding rectangles
        # overlap. do this repeatedly until you have only completely non-overlapping islands.
        consumed = set()
        prev_contours = [cv2.convexHull(c) for c in contours]
        contours = []
        for i, contour in enumerate(prev_contours):
            if not (i in consumed):
                for iother, other in enumerate(prev_contours[i+1:], i+1):
                    if rects_intersect(cv2.boundingRect(contour), cv2.boundingRect(other)):
                        contour = np.vstack([contour, other])
                        consumed.add(iother)
                contours.append(contour)
        if len(consumed) == 0:
            return contours

def contours_to_boxes(contours):
    boxes = []
    for c in contours:
        # calculate two box contours:
        # 1. axis-aligned bounding rectangle of the minimum enclosing circle
        # 2. minimum area rotated rectangle
        # use the smaller of the two (preferring #1 if they are of similar size)
        (x, y), r = cv2.minEnclosingCircle(c)
        box1 = np.intp([(x - r, y - r), (x + r, y - r), (x + r, y + r), (x - r, y + r)])
        box2 = np.intp(cv2.boxPoints(cv2.minAreaRect(c)))
        if cv2.contourArea(box1) < cv2.contourArea(box2) * 1.1:
            boxes.append(box1)
        else:
            boxes.append(box2)
    return boxes

def interactive_contour(filename):
    orig = cv2.imread(filename)
    orig = shrink(orig)

    edge = edge_img(orig)
    contours = cv2.findContours(edge, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    contours = imutils.grab_contours(contours)
    contours = merge_contours(contours)
    contours = contours_to_boxes(contours)

    cv2.createTrackbar("contour", "image", 0, len(contours) - 1, noop)

    while handle_key():
        icontour = cv2.getTrackbarPos("contour", "image")
        img = cv2.drawContours(orig.copy(), contours, icontour, (255, 0, 255), 3)
        cv2.imshow("image", img)


def get_pixel_colors(rgbimg):
    pixel_colors = rgbimg.reshape((np.shape(rgbimg)[0] * np.shape(rgbimg)[1], 3))
    norm = colors.Normalize(vmin=-1.0, vmax=1.0)
    norm.autoscale(pixel_colors)
    return norm(pixel_colors).tolist()


def plot_hsv(rgbimg):
    hsv = cv2.cvtColor(rgbimg.copy(), cv2.COLOR_BGR2HSV)

    h, s, v = cv2.split(hsv)
    fig = plt.figure()
    axis = fig.add_subplot(1, 1, 1, projection="3d")
    axis.scatter(
        h.flatten(),
        s.flatten(),
        v.flatten(),
        facecolors=get_pixel_colors(rgbimg),
        marker=".",
    )
    axis.set_xlabel("Hue")
    axis.set_ylabel("Sat")
    axis.set_zlabel("Val")
    plt.show()


def interactive_darken(filename):
    img = cv2.imread(filename)
    img = shrink(img)
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    cv2.createTrackbar("Sat", "image", 30, 255, noop)
    cv2.createTrackbar("Val", "image", 225, 255, noop)

    while handle_key():
        sat = cv2.getTrackbarPos("Sat", "image")
        val = cv2.getTrackbarPos("Val", "image")

        h, s, v = cv2.split(hsv)
        s[s > sat] = 255
        v[v < val] = 0
        hsv2 = cv2.merge((h, s, v))

        bgr = cv2.cvtColor(hsv2, cv2.COLOR_HSV2BGR)
        cv2.imshow("image", bgr)

selected_file = FILENAMES[0]
selected_tool = interactive_contour
def handle_key():
    global selected_file, selected_tool
    key = cv2.waitKey(1) & 0xff
    if key == 27: # esc
        selected_tool = None
    elif key >= ord('0') and key <= ord('9'):
        selected_file = FILENAMES[key - ord('0')]
    elif key == ord("e"):
        selected_tool = interactive_edge
    elif key == ord("d"):
        selected_tool = interactive_darken
    elif key == ord("c"):
        selected_tool = interactive_contour
    else:
        return True
    return False

def tool_loop():
    while selected_tool:
        cv2.namedWindow("image")
        selected_tool(selected_file)
        cv2.destroyWindow("image")

if __name__ == "__main__":
    tool_loop()