Boundieboxes from Pierce (OpenCV 3.0)

boundieboxes.py

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatch
import matplotlib.cm as cm
import cv2
import csv
from sklearn import cluster


def find_points(gray_img, color_img, num_points):
    corners = cv2.goodFeaturesToTrack(gray_img, num_points, 0.1, 3)
    corners = np.int0(corners)

    for corner in corners:
        x_coord, y_coord = corner.ravel()
        cv2.circle(color_img, (x_coord, y_coord), 3, 255, -1)

    corners_x = corners[:, :, 0]
    corners_y = corners[:, :, 1]
    return corners_x, corners_y, color_img


def find_orbs(image, num_points):
    # This is using the orb methodology
    orb = cv2.ORB_create(nfeatures=num_points)

    # find the key points with ORB
    kp = orb.detect(image, None)

    # compute the descriptors with ORB
    kp, des = orb.compute(image, kp)

    # draw only key points location,not size and orientation
    img2 = cv2.drawKeypoints(image, kp, outImage=None, color=(0, 255, 0), flags=0)

    # Trying to get the x,y for interesting points aka kp
    ip_x = []
    ip_y = []
    for point in kp:
        ip_x = np.append(ip_x, point.pt[0])
        ip_y = np.append(ip_y, point.pt[1])

    return ip_x, ip_y, img2


def get_boxies(point_data, unique_clusters, clusters, centroids, box_buffer, image_height, image_width):

    boxes = []
    for cluster_index in unique_clusters:
        x_points = point_data[clusters == cluster_index, 0]
        y_points = point_data[clusters == cluster_index, 1]
        x1 = int(np.percentile(x_points, 97.5) + ((box_buffer-1) * (np.percentile(x_points, 97.5) -
                                                                    np.percentile(x_points, 2.5))))
        if x1 > image_width:
            x1 = image_width
        x2 = int(np.percentile(x_points, 2.5) - ((box_buffer-1) * (np.percentile(x_points, 97.5) -
                                                                    np.percentile(x_points, 2.5))))
        if x2 < 0:
            x2 = 0
        y1 = int(np.percentile(y_points, 97.5) + ((box_buffer-1) * (np.percentile(y_points, 97.5) -
                                                                    np.percentile(y_points, 2.5))))
        if y1 > image_height:
            y1 = image_height
        y2 = int(np.percentile(y_points, 2.5) - ((box_buffer-1) * (np.percentile(y_points, 97.5) -
                                                                    np.percentile(y_points, 2.5))))
        if y2 < 0:
            y2 = 0
        box = np.array([x1, y1, x2, y2])

        boxes = np.append(boxes, box)

    boxes = boxes.reshape(len(unique_clusters), 1, 4).astype(int)
    return boxes


def plot_clusters(point_data, clusters, unique_clusters, image_height, image_width):
    colors = cm.rainbow(np.linspace(0, 1, len(unique_clusters)))

    plt.imshow(img)
    for i, c in zip(unique_clusters, colors):
        plt.scatter(point_data[clusters == i, 0], point_data[clusters == i, 1], color=c, s=0.5)

    plt.xlim(0, image_width)
    plt.ylim(0, image_height)
    plt.gca().invert_yaxis()
    plt.show()


def plot_boxies(image, box_data, image_height, image_width):

    fig, ax = plt.subplots()
    ax.imshow(image)

    for box in box_data:
        rect = mpatch.Rectangle((box[:, 2], box[:, 3]), box[:, 0]-box[:, 2], box[:, 1]-box[:, 3],
                                fill=False, edgecolor='red', linewidth=2)
        ax.add_patch(rect)

    plt.xlim(0, image_width)
    plt.ylim(0, image_height)
    plt.gca().invert_yaxis()
    plt.show()


def write_caffe_boxes(path, image, boundieboxies_output):
    header = ['filename', 'xmin', 'ymin', 'xmax', 'ymax']

    with open(path, 'w') as csv_file:
        csv_writer = csv.writer(csv_file, delimiter=',')
        csv_writer.writerow(header)
        for boxes in boundieboxies_output:
            csv_writer.writerow([image, int(boxes[:, 2]), int(boxes[:, 3]), int(boxes[:, 0]), int(boxes[:, 1])])


def boundieboxies(image, min_objects=1, max_objects=10, points_per_object=100, object_fudge_factor=1,
                  point_fudge_factor=1, cluster_initializations=10, box_buffer=1.2):

    img_height = np.shape(image)[0]
    img_width = np.shape(image)[1]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    all_boxes = []
    objects = 0
    for image_object in range(min_objects, max_objects+1):
        points = int(image_object * points_per_object * point_fudge_factor)
        n_clusters = int(image_object * object_fudge_factor)

        corners_x, corners_y, corner_pic = find_points(gray, img.copy(), points)
        ip_x, ip_y, ip_pic = find_orbs(img.copy(), points)

        x = np.append(corners_x, ip_x)
        y = np.append(corners_y, ip_y)

        data_comb = np.array([x, y]).T

        # Run K-means to find interesting localizations of interesting points
        k_means = cluster.KMeans(n_clusters=n_clusters, n_init=cluster_initializations)
        k_means.fit(data_comb)
        clusters = k_means.labels_
        centroids = k_means.cluster_centers_

        # Find unique colors
        unique_clusters = np.unique(clusters)

        boxes = get_boxies(data_comb, unique_clusters, clusters, centroids, box_buffer, img_height, img_width)

        objects += n_clusters
        all_boxes = np.append(all_boxes, boxes)

    all_boxes = all_boxes.reshape(objects, 1, 4).astype(int)
    return all_boxes


# Quick list of images to choose from
IMAGE_FILE = 'frames/3897.jpg'

# Import image with dims
img = cv2.imread(IMAGE_FILE)
img_height = np.shape(img)[0]
img_width = np.shape(img)[1]

all_box_objects = boundieboxies(img, 1, 5)

# write_caffe_boxes('/Users/pspitler3/Desktop/test.csv', IMAGE_FILE, all_box_objects)

# Plot functions
plot_boxies(img.copy(), all_box_objects, img_height, img_width)
#plot_clusters(data_comb, clusters, unique_clusters, img_height, img_width)