HelperFunctions.py

import matplotlib.image as mpimg
import numpy as np
import cv2
from skimage.feature import hog
from scipy.ndimage.measurements import label
import threading
import time

class VehicleDetection:
    def __init__(self, window_sizes, x_start_stop, y_start_stop, xy_overlap, svc, X_scalar, color_space, spatial_size,
                 hist_bins,
                 orient, pix_per_cell, cell_per_block, hog_channel, spatial_feat, hist_feat, hog_feat,
                 heatmap_threshold):

        self.window_sizes = window_sizes
        self.x_start_stop = x_start_stop
        self.curr_x_start_stop = x_start_stop
        self.y_start_stop = y_start_stop
        self.curr_y_start_stop = y_start_stop
        self.xy_overlap = xy_overlap
        self.svc = svc
        self.X_scalar = X_scalar
        self.color_space = color_space
        self.spatial_size = spatial_size
        self.hist_bins = hist_bins
        self.orient = orient
        self.pix_per_cell = pix_per_cell
        self.cell_per_block = cell_per_block
        self.hog_channel = hog_channel
        self.spatial_feat = spatial_feat
        self.hist_feat = hist_feat
        self.hog_feat = hog_feat
        self.heatmap_threshold = heatmap_threshold
        self.frame_counter = 0
        self.reset_frame_counter_every = 25




    # Define a function to return HOG features and visualization
    @staticmethod
    def get_hog_features(img, orient, pix_per_cell, cell_per_block,
                         vis=False, feature_vec=True):
        # Call with two outputs if vis==True
        if vis == True:
            features, hog_image = hog(img, orientations=orient,
                                      pixels_per_cell=(pix_per_cell, pix_per_cell),
                                      cells_per_block=(cell_per_block, cell_per_block),
                                      transform_sqrt=True,
                                      visualise=vis, feature_vector=feature_vec)
            return features, hog_image
        # Otherwise call with one output
        else:
            features = hog(img, orientations=orient,
                           pixels_per_cell=(pix_per_cell, pix_per_cell),
                           cells_per_block=(cell_per_block, cell_per_block),
                           transform_sqrt=True,
                           visualise=vis, feature_vector=feature_vec)
            return features

    # Define a function to compute binned color features
    @staticmethod
    def bin_spatial(img, size=(32, 32)):
        # Use cv2.resize().ravel() to create the feature vector
        features = cv2.resize(img, size).ravel()
        # Return the feature vector
        return features

    # Define a function to compute color histogram features
    # NEED TO CHANGE bins_range if reading .png files with mpimg!
    @staticmethod
    def color_hist(img, nbins=32, bins_range=(0, 256)):
        # Compute the histogram of the color channels separately
        channel1_hist = np.histogram(img[:, :, 0], bins=nbins, range=bins_range)
        channel2_hist = np.histogram(img[:, :, 1], bins=nbins, range=bins_range)
        channel3_hist = np.histogram(img[:, :, 2], bins=nbins, range=bins_range)
        # Concatenate the histograms into a single feature vector
        hist_features = np.concatenate((channel1_hist[0], channel2_hist[0], channel3_hist[0]))
        # Return the individual histograms, bin_centers and feature vector
        return hist_features

    # Define a function to extract features from a list of images
    # Have this function call bin_spatial() and color_hist()
    @staticmethod
    def extract_features(imgs, color_space='RGB', spatial_size=(32, 32),
                         hist_bins=32, orient=9,
                         pix_per_cell=8, cell_per_block=2, hog_channel=0,
                         spatial_feat=True, hist_feat=True, hog_feat=True):
        # Create a list to append feature vectors to
        features = []
        # Iterate through the list of images
        for file in imgs:
            file_features = []
            # Read in each one by one
            image = mpimg.imread(file)
            # apply color conversion if other than 'RGB'
            if color_space != 'RGB':
                if color_space == 'HSV':
                    feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2HSV)
                elif color_space == 'LUV':
                    feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2LUV)
                elif color_space == 'HLS':
                    feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2HLS)
                elif color_space == 'YUV':
                    feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2YUV)
                elif color_space == 'YCrCb':
                    feature_image = cv2.cvtColor(image, cv2.COLOR_RGB2YCrCb)
            else:
                feature_image = np.copy(image)

            if spatial_feat == True:
                spatial_features = VehicleDetection.bin_spatial(feature_image, size=spatial_size)
                file_features.append(spatial_features)
            if hist_feat == True:
                # Apply color_hist()
                hist_features = VehicleDetection.color_hist(feature_image, nbins=hist_bins)
                file_features.append(hist_features)
            if hog_feat == True:
                # Call get_hog_features() with vis=False, feature_vec=True
                if hog_channel == 'ALL':
                    hog_features = []
                    for channel in range(feature_image.shape[2]):
                        hog_features.append(VehicleDetection.get_hog_features(feature_image[:, :, channel],
                                                                              orient, pix_per_cell, cell_per_block,
                                                                              vis=False, feature_vec=True))
                    hog_features = np.ravel(hog_features)
                else:
                    hog_features = VehicleDetection.get_hog_features(feature_image[:, :, hog_channel], orient,
                                                                     pix_per_cell, cell_per_block, vis=False,
                                                                     feature_vec=True)
                # Append the new feature vector to the features list
                file_features.append(hog_features)
            features.append(np.concatenate(file_features))
        # Return list of feature vectors
        return features

    # Define a function that takes an image,
    # start and stop positions in both x and y,
    # window size (x and y dimensions),
    # and overlap fraction (for both x and y)
    def slide_window(self, img, x_start_stop=[None, None], y_start_stop=[None, None],
                     xy_window=(64, 64), xy_overlap=(0.5, 0.5)):
        # If x and/or y start/stop positions not defined, set to image size
        if x_start_stop[0] == None:
            x_start_stop[0] = 0
        if x_start_stop[1] == None:
            x_start_stop[1] = img.shape[1]
        if y_start_stop[0] == None:
            y_start_stop[0] = int(img.shape[0] / 2)
        if y_start_stop[1] == None:
            y_start_stop[1] = img.shape[0]
        # Compute the span of the region to be searched
        xspan = x_start_stop[1] - x_start_stop[0]
        yspan = y_start_stop[1] - y_start_stop[0]
        # Compute the number of pixels per step in x/y
        nx_pix_per_step = np.int(xy_window[0] * (1 - xy_overlap[0]))
        ny_pix_per_step = np.int(xy_window[1] * (1 - xy_overlap[1]))
        # Compute the number of windows in x/y
        nx_buffer = np.int(xy_window[0] * (xy_overlap[0]))
        ny_buffer = np.int(xy_window[1] * (xy_overlap[1]))
        nx_windows = np.int((xspan - nx_buffer) / nx_pix_per_step)
        ny_windows = np.int((yspan - ny_buffer) / ny_pix_per_step)
        # Initialize a list to append window positions to
        window_list = []
        # Loop through finding x and y window positions
        # Note: you could vectorize this step, but in practice
        # you'll be considering windows one by one with your
        # classifier, so looping makes sense
        for ys in range(ny_windows):
            for xs in range(nx_windows):
                # Calculate window position
                startx = xs * nx_pix_per_step + x_start_stop[0]
                endx = startx + xy_window[0]
                starty = ys * ny_pix_per_step + y_start_stop[0]
                endy = starty + xy_window[1]

                # Append window position to list
                window_list.append(((startx, starty), (endx, endy)))
        # Return the list of windows
        return window_list

    # Define a function to draw bounding boxes
    def draw_boxes(self, img, bboxes, color=(0, 0, 255), thick=6):
        # Make a copy of the image
        imcopy = np.copy(img)
        # Iterate through the bounding boxes
        for bbox in bboxes:
            # Draw a rectangle given bbox coordinates
            cv2.rectangle(imcopy, bbox[0], bbox[1], color, thick)
        # Return the image copy with boxes drawn
        return imcopy

    def add_heat(self, heatmap, bbox_list):
        # Iterate through list of bboxes
        for box in bbox_list:
            # Add += 1 for all pixels inside each bbox
            # Assuming each "box" takes the form ((x1, y1), (x2, y2))
            heatmap[box[0][1]:box[1][1], box[0][0]:box[1][0]] += 1

        # Return updated heatmap
        return heatmap  # Iterate through list of bboxes

    def apply_threshold(self, heatmap, threshold):
        # Zero out pixels below the threshold
        heatmap[heatmap <= threshold] = 0
        # Return thresholded map
        return heatmap

    def draw_labeled_bboxes(self, img, labels):

        # cv2.rectangle(img, (self.curr_x_start_stop[0], self.curr_y_start_stop[0]),
        #               (self.curr_x_start_stop[1], self.curr_y_start_stop[1]), (0, 255, 0), 4)

        # If we don't have anything to draw just return
        if labels[1] == 0:
            return img

        # Initializing curr_min_x to be the last max x to find something that is small
        curr_min_x = self.x_start_stop[1]
        curr_min_y = self.y_start_stop[1]
        curr_max_x = self.x_start_stop[0]
        curr_max_y = self.y_start_stop[0]

        # Iterate through all detected cars
        for car_number in range(1, labels[1] + 1):
            # Find pixels with each car_number label value
            nonzero = (labels[0] == car_number).nonzero()
            # Identify x and y values of those pixels
            nonzeroy = np.array(nonzero[0])
            nonzerox = np.array(nonzero[1])
            # Define a bounding box based on min/max x and y
            min_x = np.min(nonzerox)
            if min_x < curr_min_x:
                curr_min_x = min_x
            min_y = np.min(nonzeroy)
            if min_y < curr_min_y:
                curr_min_y = min_y

            max_x = np.max(nonzerox)
            if max_x > curr_max_x:
                curr_max_x = max_x
            max_y = np.max(nonzeroy)
            if max_y > curr_max_y:
                curr_max_y = max_y

            bbox = ((min_x, min_y), (max_x, max_y))

            # print(bbox)

            # Draw the box on the image
            cv2.rectangle(img, bbox[0], bbox[1], (0, 0, 255), 4)

        self.curr_x_start_stop = [curr_min_x - int(curr_min_x * .1), curr_max_x + int(curr_max_x * .1)]
        # Guards so that we do not exceed the global search window size
        if self.curr_x_start_stop[0] < self.x_start_stop[0]:
            self.curr_x_start_stop[0] = self.x_start_stop[0]
        if self.curr_x_start_stop[1] > self.x_start_stop[1]:
            self.curr_x_start_stop[1] = self.x_start_stop[1]

        self.curr_y_start_stop = [curr_min_y - int(curr_min_y * .1), curr_max_y + int(curr_max_y * .1)]
        # Guards so that we do not exceed the global search window size
        if self.curr_y_start_stop[0] < self.y_start_stop[0]:
            self.curr_y_start_stop[0] = self.y_start_stop[0]
        if self.curr_y_start_stop[1] > self.y_start_stop[1]:
            self.curr_y_start_stop[1] = self.y_start_stop[1]

        # Return the image
        return img

    @staticmethod
    def convert_color(img, conv='RGB2YCrCb'):
        if conv == 'RGB2YCrCb':
            return cv2.cvtColor(img, cv2.COLOR_RGB2YCrCb)
        if conv == 'BGR2YCrCb':
            return cv2.cvtColor(img, cv2.COLOR_BGR2YCrCb)
        if conv == 'RGB2LUV':
            return cv2.cvtColor(img, cv2.COLOR_RGB2LUV)

    # Define a single function that can extract features using hog sub-sampling and make predictions

    def process_frame(self, image):
        draw_image = np.copy(image)

        # Uncomment the following line if you extracted training
        # data from .png images (scaled 0 to 1 by mpimg) and the
        # image you are searching is a .jpg (scaled 0 to 255)

        image = image.astype(np.float32) / 255

        hot_windows = []

        # Reset search window every x frames
        if self.frame_counter % self.reset_frame_counter_every == 0:
            # print("Resetting Search Window Size")
            self.heat = np.zeros_like(image[:, :, 0]).astype(np.float)
            self.curr_x_start_stop = self.x_start_stop
            self.curr_y_start_stop = self.y_start_stop

        self.frame_counter += 1

        threads = []
        # fork a thread for each window size
        for size in self.window_sizes:
            # windows = self.slide_window(image, x_start_stop=self.curr_x_start_stop, y_start_stop=self.curr_y_start_stop,
            #                             xy_window=(size, size), xy_overlap=self.xy_overlap)
            #
            # hot_windows.extend(
            #     self.search_windows(image, windows, self.svc, self.X_scalar, color_space=self.color_space,
            #                         spatial_size=self.spatial_size, hist_bins=self.hist_bins,
            #                         orient=self.orient, pix_per_cell=self.pix_per_cell,
            #                         cell_per_block=self.cell_per_block,
            #                         hog_channel=self.hog_channel, spatial_feat=self.spatial_feat,
            #                         hist_feat=self.hist_feat, hog_feat=self.hog_feat))

            thread = FindVehicles(image, self.x_start_stop, self.curr_y_start_stop, size / 64, self.svc,
                                  self.X_scalar, self.orient, self.pix_per_cell, self.cell_per_block, self.spatial_size,
                                  self.hist_bins)
            thread.name = "Scale:{}".format(size/64)
            threads.append(thread)
            thread.start()
        # wait for the threads until they finish
        for t in threads:
            t.join()
        # Gather the hot windows (boxes) from the different threads
        for t in threads:
            hot_windows.extend(t.hot_windows)

        # print(hot_windows)

        # Add heat to each hot boxes in box list
        self.heat = self.add_heat(self.heat, hot_windows)

        # Apply threshold to help remove false positives
        self.heat = self.apply_threshold(self.heat, self.heatmap_threshold)

        # Visualize the heatmap when displaying
        # heatmap = np.clip(self.heat, 0, 255)

        # Find final boxes from heatmap using label function
        labels = label(self.heat)
        window_img = self.draw_labeled_bboxes(draw_image, labels)
        # window_img = self.draw_boxes(draw_image, hot_windows, color=(0, 0, 255), thick=4)

        # window_img = cv2.cvtColor(window_img, cv2.COLOR_BGR2RGB)

        return window_img

    # Define a function to extract features from a single image window
    # This function is very similar to extract_features()
    # just for a single image rather than list of images
    def single_img_features(self, img, color_space='RGB', spatial_size=(32, 32),
                            hist_bins=32, orient=9,
                            pix_per_cell=8, cell_per_block=2, hog_channel=0,
                            spatial_feat=True, hist_feat=True, hog_feat=True):
        # 1) Define an empty list to receive features
        img_features = []
        # 2) Apply color conversion if other than 'RGB'
        if color_space != 'RGB':
            if color_space == 'HSV':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
            elif color_space == 'LUV':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2LUV)
            elif color_space == 'HLS':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2HLS)
            elif color_space == 'YUV':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2YUV)
            elif color_space == 'YCrCb':
                feature_image = cv2.cvtColor(img, cv2.COLOR_RGB2YCrCb)
        else:
            feature_image = np.copy(img)
        # 3) Compute spatial features if flag is set
        if spatial_feat is True:
            spatial_features = self.bin_spatial(feature_image, size=spatial_size)
            # 4) Append features to list
            img_features.append(spatial_features)
        # 5) Compute histogram features if flag is set
        if hist_feat == True:
            hist_features = self.color_hist(feature_image, nbins=hist_bins)
            # 6) Append features to list
            img_features.append(hist_features)
        # 7) Compute HOG features if flag is set
        if hog_feat == True:
            if hog_channel == 'ALL':
                hog_features = []
                for channel in range(feature_image.shape[2]):
                    hog_features.extend(self.get_hog_features(feature_image[:, :, channel],
                                                              orient, pix_per_cell, cell_per_block,
                                                              vis=False, feature_vec=True))
            else:
                hog_features = self.get_hog_features(feature_image[:, :, hog_channel], orient,
                                                     pix_per_cell, cell_per_block, vis=False, feature_vec=True)
            # 8) Append features to list
            img_features.append(hog_features)

        # 9) Return concatenated array of features
        return np.concatenate(img_features)

    # Define a function you will pass an image
    # and the list of windows to be searched (output of slide_windows())
    def search_windows(self, img, windows, clf, scaler, color_space='RGB',
                       spatial_size=(32, 32), hist_bins=32,
                       hist_range=(0, 256), orient=9,
                       pix_per_cell=8, cell_per_block=2,
                       hog_channel=0, spatial_feat=True,
                       hist_feat=True, hog_feat=True):
        # 1) Create an empty list to receive positive detection windows
        on_windows = []
        # 2) Iterate over all windows in the list
        for window in windows:
            # 3) Extract the test window from original image
            test_img = cv2.resize(img[window[0][1]:window[1][1], window[0][0]:window[1][0]], (64, 64))
            # 4) Extract features for that window using single_img_features()
            features = self.single_img_features(test_img, color_space=color_space,
                                                spatial_size=spatial_size, hist_bins=hist_bins,
                                                orient=orient, pix_per_cell=pix_per_cell,
                                                cell_per_block=cell_per_block,
                                                hog_channel=hog_channel, spatial_feat=spatial_feat,
                                                hist_feat=hist_feat, hog_feat=hog_feat)
            # 5) Scale extracted features to be fed to classifier
            test_features = scaler.transform(np.array(features).reshape(1, -1))
            # 6) Predict using your classifier
            prediction = clf.predict(test_features)
            # 7) If positive (prediction == 1) then save the window
            if prediction == 1:
                on_windows.append(window)
        # 8) Return windows for positive detections
        return on_windows


class FindVehicles(threading.Thread):
    def __init__(self, img, x_start_stop, y_start_stop, scale, svc, X_scaler, orient, pix_per_cell, cell_per_block,
                 spatial_size, hist_bins, cells_per_step=2):
        # draw_img = np.copy(img)
        # img = img.astype(np.float32) / 255

        threading.Thread.__init__(self)

        self.img = img
        self.x_start_stop = x_start_stop
        self.y_start_stop = y_start_stop
        self.scale = scale
        self.svc = svc
        self.X_scaler = X_scaler
        self.orient = orient
        self.pix_per_cell = pix_per_cell
        self.cell_per_block = cell_per_block
        self.spatial_siz = spatial_size
        self.hist_bins = hist_bins
        self.cells_per_step = cells_per_step
        self.hot_windows = []

    def run(self):

        t_start = time.time()
        img_tosearch = self.img[self.y_start_stop[0]:self.y_start_stop[1], self.x_start_stop[0]:self.x_start_stop[1], :]
        ctrans_tosearch = VehicleDetection.convert_color(img_tosearch, conv='RGB2YCrCb')
        if self.scale != 1:
            imshape = ctrans_tosearch.shape
            ctrans_tosearch = cv2.resize(ctrans_tosearch,
                                         (np.int(imshape[1] / self.scale), np.int(imshape[0] / self.scale)))

        ch1 = ctrans_tosearch[:, :, 0]
        ch2 = ctrans_tosearch[:, :, 1]
        ch3 = ctrans_tosearch[:, :, 2]

        # Define blocks and steps as above
        nxblocks = (ch1.shape[1] // self.pix_per_cell) - self.cell_per_block + 1
        nyblocks = (ch1.shape[0] // self.pix_per_cell) - self.cell_per_block + 1
        nfeat_per_block = self.orient * self.cell_per_block ** 2

        # 64 was the orginal sampling rate, with 8 cells and 8 pix per cell
        window = 64
        nblocks_per_window = (window // self.pix_per_cell) - self.cell_per_block + 1
        # Instead of overlap, define how many cells to step, cells_per_step
        nxsteps = (nxblocks - nblocks_per_window) // self.cells_per_step
        nysteps = (nyblocks - nblocks_per_window) // self.cells_per_step

        # Compute individual channel HOG features for the entire image
        hog1 = VehicleDetection.get_hog_features(ch1, self.orient, self.pix_per_cell, self.cell_per_block,
                                                 feature_vec=False)
        hog2 = VehicleDetection.get_hog_features(ch2, self.orient, self.pix_per_cell, self.cell_per_block,
                                                 feature_vec=False)
        hog3 = VehicleDetection.get_hog_features(ch3, self.orient, self.pix_per_cell, self.cell_per_block,
                                                 feature_vec=False)

        for xb in range(nxsteps):
            for yb in range(nysteps):
                ypos = yb * self.cells_per_step
                xpos = xb * self.cells_per_step
                # Extract HOG for this patch
                hog_feat1 = hog1[ypos:ypos + nblocks_per_window, xpos:xpos + nblocks_per_window].ravel()
                hog_feat2 = hog2[ypos:ypos + nblocks_per_window, xpos:xpos + nblocks_per_window].ravel()
                hog_feat3 = hog3[ypos:ypos + nblocks_per_window, xpos:xpos + nblocks_per_window].ravel()
                hog_features = np.hstack((hog_feat1, hog_feat2, hog_feat3))

                xleft = xpos * self.pix_per_cell
                ytop = ypos * self.pix_per_cell

                # Extract the image patch
                subimg = cv2.resize(ctrans_tosearch[ytop:ytop + window, xleft:xleft + window], (64, 64))

                # Get color features
                # spatial_features = self.bin_spatial(subimg, size=spatial_size)
                hist_features = VehicleDetection.color_hist(subimg, nbins=self.hist_bins)

                # Scale features and make a prediction
                test_features = self.X_scaler.transform(
                    np.concatenate((hist_features, hog_features)).reshape(1, -1))
                # test_features = X_scaler.transform(np.hstack((shape_feat, hist_feat)).reshape(1, -1))
                test_prediction = self.svc.predict(test_features)

                if test_prediction == 1:
                    xbox_left = np.int(xleft * self.scale)
                    ytop_draw = np.int(ytop * self.scale)
                    win_draw = np.int(window * self.scale)
                    self.hot_windows.append(
                        ((xbox_left, ytop_draw + self.y_start_stop[0]),
                         (xbox_left + win_draw, ytop_draw + win_draw + self.y_start_stop[0])))

                    # cv2.rectangle(draw_img, (xbox_left, ytop_draw + ystart),
                    #               (xbox_left + win_draw, ytop_draw + win_draw + ystart), (0, 0, 255), 6)

        # print("Thread:{} took:{} secs".format(self._name, time.time() - t_start))