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


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.namedWindow("image")
    cv2.createTrackbar("threshold1", "image", 5, 50, noop)
    cv2.createTrackbar("threshold2", "image", 20, 50, noop)

    while 1:
        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)
        if cv2.waitKey(0) < 0:
            break
    cv2.destroyWindow("image")


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 autocrop_contour(orig, contour):
    approx = approximate_contour(contour)
    if len(approx) == 4:
        # It's a rectangle - perspective-correct it
        pass
    else:
        # It's a disc - crop around it
        pass


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 = sorted(contours, key=cv2.contourArea, reverse=True)
    contours = [approximate_contour(c) for c in contours]

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

    while True:
        icontour = cv2.getTrackbarPos("contour", "image")
        img = cv2.drawContours(orig.copy(), contours, icontour, (255, 0, 255), 3)
        cv2.imshow("image", img)
        if cv2.waitKey() < 0:
            break


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.namedWindow("image")
    cv2.createTrackbar("Sat", "image", 30, 255, noop)
    cv2.createTrackbar("Val", "image", 225, 255, noop)

    while True:
        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)
        if cv2.waitKey() < 0:
            break


if __name__ == "__main__":
    interactive_edge("testimg3.jpg")