sign_classifier.py

import numpy as np
import joblib
import os

MODEL_ROOT = "/home/jetbot/jetbot/notebooks/UoA/models"

class SignClassifier():
    """Responsible for classifying traffic signs"""
    
    def __init__(self, model_path):
        """Loads trained model"""
        self.clf = joblib.load(os.path.join('./models', model_path))  # Assume SK learn joblib export.

    
    def minmax_scale(self, arr):
        """Performs min-mx normalisation on input array"""
        min_value = np.min(arr)
        max_value = np.max(arr)

        return (arr - min_value) / (max_value - min_value)
    
    
    def bgr_to_cmyk(self, bgr):
        """Converts BGR8 into the CMYK colour space.
        
        Adapted from stackoverflow.com/questions/69955216
        """
        bgrdash = bgr.astype(np.float)/255.

        K = 1 - np.max(bgrdash, axis=2)
        C = (1-bgrdash[:, :, 2] - K)/(1-K)
        M = (1-bgrdash[:, :, 1] - K)/(1-K)
        Y = (1-bgrdash[:, :, 0] - K)/(1-K)

        return (np.dstack((C,M,Y,K)) * 255).astype(np.uint8)


    def get_feature_vector(self, snapshot):
        """Creates feature vector for current camera state"""
        return self.minmax_scale(self.bgr_to_cmyk(snapshot).flatten())
    

    def get_prediction(self, snapshot):
        """Makes a prediction with the current camera state. 
        Image should be in the BGR colour space and should be a 3d numpy array of uint8 data types.
        """
        return self.clf.predict_proba([self.get_feature_vector(snapshot)])